HackerRanK: Tree Pruning

HackerRanK: Tree Pruning

Given a tree with N vertices numbered from 1 to N. The vertex 1 is the root of the tree. Each vertex is assigned with an integer weight. A remove operation can remove sub-tree rooted at an arbitrary vertex u. You must use at most K remove operations (possibly zero) so that the total weight of the remaining vertices is largest.

Input Format
The first line contains two integers N,K.
The second line contains N integers, the ith integer is the weight of the ith vertex.
Each of the next N−1 lines contains a pair of integers u and v, which represents an edge from vertex u to vertex v.

There is an interesting observation on the tree is that when you consider the nodes of the tree in the DFS order (the order that the node is visited when we DFS from the root) you will see that each nodes u will be followed by all the nodes in the sub-tree rooted at u. So after do the DFS and get some necessary information like the order the nodes is visited and the size of the sub-tree rooted at each node we can forget about the tree structure and focus only on the sequence of the nodes in the DFS order.

Let A[i] (1 <= i <= N) is the ith node we visited in the DFS process; W[i] is weight of the node i, S[i] is the number of nodes in the sub-tree rooted at node i. A sub-tree will corresponding to a segment of contiguous elements on array A such that the first elements is some node u and following u in the segment is exactly S[u] - 1 nodes more. We will have to remove at most K such segments and maximize the total weight of the remaining nodes.

Let F[] is the maximum total weight of the nodes in the range A[1..i] if we removed at most j segments in this range. Our needed result will be F[N][K] and a simple dynamic programming algorithm will help us calculate F. There are only two scenario:

• if you do not remove A[i], we have: F[i][j] = max(F[i][j], F[i - 1][j] + W[A[i]])
• if you remove the segment started at A[i] then we have: F[i + S[A[i]]][j + 1] = max(F[i + S[A[i]]][j + 1], F[i][j]). This gives you the dynamic programming formula. The initiation is left for you to decides yourself, basically you can just set F[0][i] = 0 for all i and set a large enough value for all the other F[i][j]. The complexity of this approach is O(NK).