HackerRank Weekly #11 Problem 3

I worked on this problem a couple of weeks ago during HackerRank's Weekly Challenge #11 and failed to solve it during the contest. Following the contest an editorial of the problem was made available. However it just had a broad description of the correct approach and, despite the correct code, the narrative contained some incorrect details regarding the specifics, which is partly the reason I wanted to write this post. I also wanted to have a proper writeup of this question so that I have some solid documentation on a not-so-easy problem solution of which can incorrectly be considered as greedy. Which is why, after a brief description of the question, I will explain the incorrect greedy approach I initially thought would work. Then a demonstration of how it fails will follow. Finally I will state the dynamic programming solution, including the code, and hopefully with no mistakes.

Problem Statement

The statement can be found here. We are given a tree with N nodes with weights on each node. We are then asked to prune it not more than K times, to get a tree with maximum possible weight. Obviously the wieghts of nodes, thus subtrees, can be negative. The bound on the N and K are 100000 and 200, respectively.

Why Not Greedy?

Once I read the statement the thought of pruning the subtree with the least weight first seemed like the rational thing to do. Figured I would do that until I either made K cuts or no subtree has negative weight. All needed to be done to get the weights was to traverse the tree with, say, depth-first search and put the weights of the subtrees into a priority queue. Then the minimum could easily be retrieved. Also, upon deletion of a subtree, all its parents could be checked in O(logN) time and a more up-to-date version of their subtrees' weights could be put into the priority queue. As we delete subtrees we could update the weight of the overall tree in a variable which we could, in the end, print out.

The idea seemed correct during the contest and I tried quite a few times to solve the problem with this approach, considering my failure to be due to implementing the solution incorrectly. It was only after the contest, as I was analyzing one of the smaller test cases I saw the problem with the greedy method. I will cut it short and demonstrate a case where greedy approach fails.

6 2  
418 -874 -850 -257 860 -424  
1 2  
1 3  
3 4  
4 5  
4 6  

In the case given above, since the overall weight of the tree is -1127, which also happens to be the lowest value among all the subtree weights, a greedy algorithm tries to cut the whole tree in its first pruning attempt, and gets an empty tree with weight 0. Yet there exists a better way, which consists of cutting down nodes 2 and 3, thereby getting a tree consisting only of node 1, which has a weight of 418. If we were only allowed to do a single cut, the greedy solution was going to be correct. Thus, doing the most optimal cut at each instance does not yield the best result.

The important thing to notice here is that greedy solutions go step by step, by making the optimal move on top of the last (i.e. then-optimal) step. Dynamic programming, on the other hand, takes into account the actions that seem sub-optimal on its own, and manages to yield optimal solutions which are in fact combination of those individually sub-optimal moves. This can also be seen when time complexities of the solutions are compared. The greedy solution I mentioned has O(Nlog(N) + Klog(N)) complexity while the dynamic programming solution has O(N + NK) complexity. Omitting the traversal of the tree, the complexities show that greedy just makes K consecutive moves while dynamic programming solution tinkers with different combinations of moves.

Dynamic Programming Solution

This solution requires the depth-first traversal of the tree as well, to record the size of each subtree in an array S, as well as the order in which they are visited in an array A. It is important to note that, after the traversal, for the ith node that was visited, the inclusive range of nodes from A[i] to A[i + S[A[i]] - 1] contains the subtree, root of which is A[i].

Dynamic programming approach maintains a two-dimensional array F. The value maintained in F[i][j] is the maximum tree weight that can be obtained in the inclusive range of nodes A[1] to A[i], making at most j cuts.

For any node A[i] discovered during traversal, there are two options. We either cut off(i.e. prune) that node, or we don't. Among these two options we need to pick the one that yields a greater tree weight. Let's consider the value of F[i][j] for the latter case. If we are not to prune that node, we need only to add the weight of node A[i] to F[i-1][j], which should be clear if you consider the definition of F[i][j] made in the previous paragraph.

The second case is when we prune that node. Note that we are thinking of including in terms of individual nodes, and not subtrees, whereas we prune the whole subtree. Thus, choosing to ignore the subtree rooted at A[i] means ignoring all the nodes in that subtree, all S[A[i]] of them. So, cutting the subtree rooted at A[i] means updating F[i + S[A[i] - 1][j+1] to reflect the value of F[i-1][j], which is the best weight for the first i-1 nodes. Notice that we updated j to j+1. After all, pruning means making another cut.

As can be seen below in the code, while investigating F[i][j], this approach may update F[a][b] such that a > i or b > j. Also, any cell in F may be updated multiple times. So, how we need to initialize and update F carefully. Updating any cell F[i][j] using max() function and having F[i][j] as one of its parameters is reasonable, as F[i][j] should store the maximum value. To initialize F[i][j], for all values with i and j greater than 0, LLONG_MAX in climits can be used. This ensures that updating by max() causes no issues. Also, we need to assign 0 to all F[0][j]. For the dynamic computation of the cases F[i][j], we need to start iterating from (j == 0), as values of that column may affect the values of the ones to its right. In addition to that, the dynamic computation is executed in a column-first manner. (i.e. the loop for j is the outer loop) Finally, I did not optimize the computation to use one dimensional array. I believe the code is easier to read it this way. However, if you're interested, such a space-optimized version can be found on the code provided with the editorial of the problem.

Code

#include <cmath>
#include <cstdio>
#include <vector>
#include <climits>
#include <algorithm>
#define NLIMIT 100001
#define KLIMIT 201

using namespace std;

vector<int> adjList[NLIMIT];  
bool visited[NLIMIT];

int subtreeSize[NLIMIT], visitOrder[NLIMIT];  
long long weight[NLIMIT], dp[NLIMIT][KLIMIT];  
int nextToVisit;

void dfs(int u, int parent)  
{
    int discoveryTime = nextToVisit;
    visitOrder[nextToVisit++] = u;
    visited[u] = true;

    for(int i=adjList[u].size() - 1; i >= 0; --i)
    {
        if(!visited[adjList[u][i]])
        {
            dfs(adjList[u][i], u);
        }
    }

    subtreeSize[u] = nextToVisit - discoveryTime;
}

int main()  
{
    int a, b, n, k;
    long long maxWeight = 0;
    scanf("%d%d", &n, &k);
    for(int i=1; i <= n; ++i)
    {
        scanf("%lld", weight + i);
        maxWeight += weight[i];
    }

    for(int i=1; i < n; ++i)
    {
        scanf("%d%d", &a, &b);
        adjList[a].push_back(b);
        adjList[b].push_back(a);
    }

    for(int i=0; i <= k; ++i)
    {
        dp[0][i] = 0;
    }
    for(int i=1; i <= n; ++i)
    {
        for(int j=1; j <= k; ++j)
        {
            dp[i][j] = LLONG_MIN;
        }
    }
    nextToVisit = 1;
    dfs(1, 0);
    for(int i=1; i <= n; ++i)
    {
        dp[i][0] = weight[visitOrder[i]] + dp[i-1][0];
    }


    for(int j=0; j <= k; ++j)
    {
        for(int i=1; i <= n; ++i)
        {
            dp[ i + subtreeSize[ visitOrder[i] ] - 1][j + 1] = max(dp[ i + subtreeSize[ visitOrder[i] ] -1][j + 1],
                                                               dp[i-1][j]);
            dp[i][j] = max(dp[i-1][j] + weight[visitOrder[i]], dp[i][j]);
        }
    }

    printf("%lld\n", dp[n][k]);

    return 0;
}

Conclusion

It still is relatively challenging to notice whether a greedy approach would work for a problem during a programming contest. A formal proof that greedy approach would fail is not that easy to manage, as there is limited time to work on it during the contest. Greedy ideas usually seem intuitive to one, which is why they are hard to prove wrong. I guess it all comes down to noticing a case where greedy approach would be incorrect, or not optimal. In my experience, the best way to be able to think of or see such cases instantly is to get a lot of practice with greedy problems as well as dynamic programming ones. Not too different from the only way to get better at competitive programming, I guess.