Maximum Weight Subset (Codeforces 1249F)
Tags: DP, trees

You are given a tree, which consists of $$$n$$$ vertices. Recall that a tree is a connected undirected graph without cycles.   Example of a tree. Vertices are numbered from $$$1$$$ to $$$n$$$. All vertices have weights, the weight of the vertex $$$v$$$ is $$$a_v$$$.Recall that the distance between two vertices in the tree is the number of edges on a simple path between them.Your task is to find the subset of vertices with the maximum total weight (the weight of the subset is the sum of weights of all vertices in it) such that there is no pair of vertices with the distance $$$k$$$ or less between them in this subset.

Input: The first line of the input contains two integers $$$n$$$ and $$$k$$$ ($$$1 \le n, k \le 200$$$) — the number of vertices in the tree and the distance restriction, respectively.The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$1 \le a_i \le 10^5$$$), where $$$a_i$$$ is the weight of the vertex $$$i$$$.The next $$$n - 1$$$ lines contain edges of the tree. Edge $$$i$$$ is denoted by two integers $$$u_i$$$ and $$$v_i$$$ — the labels of vertices it connects ($$$1 \le u_i, v_i \le n$$$, $$$u_i \ne v_i$$$).It is guaranteed that the given edges form a tree.

Output: Print one integer — the maximum total weight of the subset in which all pairs of vertices have distance more than $$$k$$$.


Sample Inputs:
5 1
1 2 3 4 5
1 2
2 3
3 4
3 5

7 2
2 1 2 1 2 1 1
6 4
1 5
3 1
2 3
7 5
7 4

Sample Outputs:
11

4

Solution (Python):
import sys

def main():
    input = sys.stdin.readline
    (N, K) = map(int, input().split())
    (*A,) = map(int, input().split())
    G = [[] for _ in range(N)]
    for _ in range(N - 1):
        (u, v) = map(int, input().split())
        (u, v) = (u - 1, v - 1)
        G[u].append(v)
        G[v].append(u)
    K += 1
    dp = [[0] * (K + 1) for _ in range(N)]
    P = [-1] * N
    P[0] = 0
    stk = [0]
    while stk:
        v = stk.pop()
        if v >= 0:
            stk.append(~v)
            for nv in G[v]:
                if P[nv] >= 0:
                    continue
                P[nv] = v
                stk.append(nv)
            continue
        v = ~v
        p = P[v]
        d0 = 0
        C = []
        for nv in G[v]:
            if nv == p:
                continue
            d0 += dp[nv][K - 1]
            C.append(nv)
        dp[v][0] = d0 + A[v]
        if len(C) >= 2:
            for d in range(1, K):
                d1 = d
                d2 = max(d, K - d1)
                if d1 > d2:
                    (d1, d2) = (d2, d1)
                (d1, d2) = (d1 - 1, d2 - 1)
                S = sum((dp[c][d2] for c in C))
                dp[v][d] = max((S - dp[c][d2] + dp[c][d1] for c in C))
        elif len(C) == 1:
            c = C[0]
            for d in range(1, K):
                dp[v][d] = dp[c][d - 1]
        for d in range(K - 1, -1, -1):
            dp[v][d] = max(dp[v][d], dp[v][d + 1])
    print(dp[0][0])
if __name__ == '__main__':
    main()

Solution (Python):
from collections import deque
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def bfs(s):
    q = deque()
    q.append(s)
    dist = [-1] * (n + 1)
    dist[s] = 0
    p = []
    child = [[] for _ in range(n + 1)]
    while q:
        i = q.popleft()
        p.append(i)
        di = dist[i]
        for j in G[i]:
            if dist[j] == -1:
                dist[j] = di + 1
                child[i].append(j)
                q.append(j)
    return (dist, p, child)
(n, k) = map(int, input().split())
k += 1
a = [0] + list(map(int, input().split()))
G = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    G[u].append(v)
    G[v].append(u)
(dist, p, child) = bfs(1)
dp = [[0] * (2 * n + 1) for _ in range(n + 1)]
while p:
    i = p.pop()
    dpi = dp[i]
    di = dist[i]
    c = child[i]
    for j in range(di + 1, n + 1):
        x = max(j, k - j + 2 * di)
        for u in c:
            s = 0
            for v in c:
                if u ^ v:
                    s += dp[v][x]
                else:
                    s += dp[v][j]
            dpi[j] = max(dpi[j], s)
    s = 0
    for u in c:
        s += dp[u][di + k]
    dpi[di] = a[i] + s
    for j in range(2 * n - 1, -1, -1):
        dpi[j] = max(dpi[j], dpi[j + 1])
ans = 0
for dp0 in dp:
    ans = max(ans, max(dp0))
print(ans)

Solution (Python):
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
g = {}

def dfs(v, p=-1):
    c = [dfs(child, v) for child in g.get(v, set()) - {p}]
    c.sort(key=len, reverse=True)
    r = []
    i = 0
    while c:
        if i >= len(c[-1]):
            c.pop()
        else:
            o = max(i, k - i - 1)
            s = q = 0
            for x in c:
                if len(x) <= o:
                    q = max(q, x[i])
                else:
                    s += x[o]
                    q = max(q, x[i] - x[o])
            r.append(q + s)
            i += 1
    r.append(0)
    for i in range(len(r) - 1, 0, -1):
        r[i - 1] = max(r[i - 1], r[i])
    while len(r) > 1 and r[-2] == 0:
        r.pop()
    o = (r[k] if k < len(r) else 0) + a[v]
    r.insert(0, max(o, r[0]))
    return r
for _ in range(1, n):
    (u, v) = map(lambda x: int(x) - 1, input().split())
    g.setdefault(u, set()).add(v)
    g.setdefault(v, set()).add(u)
print(dfs(0)[0])