Maximum Weight Subset
Contest: Codeforces Round 595 (Div. 3)
Difficulty rating: 2200
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 format:
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 format:
Print one integer — the maximum total weight of the subset in which all pairs of vertices have distance more than $$$k$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

Example input:
5 1
1 2 3 4 5
1 2
2 3
3 4
3 5

Example output:
11

Example input:
7 2
2 1 2 1 2 1 1
6 4
1 5
3 1
2 3
7 5
7 4

Example output:
4

Editorial:
Let's solve this problem using dynamic programming on a tree. Suppose the tree is rooted and the root of the tree is $$$1$$$. Also, let's increase $$$k$$$ to find the subset in which any pair of vertices had distance $$$k$$$ or greater instead of $$$k+1$$$ or greater.
Let $$$dp_{v, dep}$$$ be the maximum total weight of the subset in the subtree of $$$v$$$ if the vertex with the minimum depth we took has depth at least $$$dep$$$. Then the answer is $$$dp_{1, 0}$$$.
Firstly, let's calculate this dynamic programming for all children of $$$v$$$. Then we are ready to calculate all $$$dp_{v, dep}$$$ for all $$$dep$$$ from $$$0$$$ to $$$n$$$.
Let the current depth be $$$dep$$$, then there are two cases: if $$$dep = 0$$$ then $$$dp_{v, dep} = a_v + \sum\limits_{to \in children(v)} dp_{to, max(0, k - dep - 1)}$$$. Otherwise, let's iterate over all children of $$$v$$$ and let $$$to$$$ be such child of $$$v$$$ that the vertex with the minimum depth we took is in the subtree of $$$to$$$. Then $$$dp_{v, dep} = max(dp_{v, dep}, dp_{to, dep - 1} + \sum\limits_{other \in children(v) \backslash \{to\}} dp_{other, max(dep - 1, k - dep - 1)}$$$.
After we calculated all values of $$$dp_{v, dep}$$$ for the vertex $$$v$$$, we can notice that this is not what we wanted. The current value of $$$dp_{v, dep}$$$ means the maximum total weight of the subset in the subtree of $$$v$$$ if the vertex with the minimum depth we took has depth exactly $$$dep$$$. To fix this, let's push $$$dp_{v, dep}$$$ to $$$dp_{v, dep - 1}$$$ for all depths from $$$n$$$ to $$$1$$$.
Time complexity: $$$O(n^3)$$$ but it can be easily optimized to $$$O(n^2)$$$ using some prefix and suffix maximums.
You can ask, why this is $$$O(n^3)$$$ but not $$$O(n^4)$$$ because we iterating over all vertices, then over all possible depths, and then over children of the vertex, and again over children of the vertex. But in fact, this is $$$O(n^3)$$$ because if we change the order of multiplication, we can see that we are iterating over pairs (parent, child), then over children and possible depths, and the number of such pairs is $$$O(n)$$$, so the complexity is $$$O(n^3)$$$.