提出 #68015518

---

ソースコード 拡げる

import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')
import sys
sys.setrecursionlimit(2 * 10**7)
input = sys.stdin.readline
n, k = map(int, input().split())
a = list(map(int, input().split()))
g = [[] for _ in range(n)]
for _ in range(n - 1):
    u, v = map(int, input().split())
    g[u - 1].append(v - 1)
    g[v - 1].append(u - 1)
INF = 10**18
def f(to, fr):
    d0 = [-INF] * (k + 1)
    d1 = [-INF] * (k + 1)
    d2 = [-INF] * (k + 1)
    d0[0] = 0
    for x in g[to]:
        if x == fr:
            continue
        new_d0 = [-INF] * (k + 1)
        new_d1 = [-INF] * (k + 1)
        new_d2 = [-INF] * (k + 1)
        di0, di1 = f(x, to)
        for i in range(k + 1):
            for j in range(k + 1):
                if i + j <= k:
                    new_d0[i + j] = max(new_d0[i + j], d0[i] + di0[j])
                    new_d1[i + j] = max(new_d1[i + j], d0[i] + di1[j] + a[to])
                    new_d1[i + j] = max(new_d1[i + j], d1[i] + di0[j])
                    new_d2[i + j] = max(new_d2[i + j], d2[i] + di0[j])
                    # new_d1[i + j] = max(new_d1[i + j], d1[i] + di1[j])
                    new_d2[i + j] = max(new_d2[i + j], d2[i] + di1[j])
                if 0 <= i + j - 1 <= k:
                    new_d2[i + j - 1] = max(new_d2[i + j - 1], d1[i] + di1[j])
        d0, d1, d2 = new_d0, new_d1, new_d2
    for i in range(1, k + 1):
        d1[i] = max(d1[i], d0[i - 1] + a[to])
    for i in range(k + 1):
        d0[i] = max(d0[i], d1[i], d2[i])
    return d0, d1
d0, d1 = f(0, -1)
print(max(d0))

提出情報

ジャッジ結果