In this tutorial we will be discussing dynamic programming on trees, a very popular algorithmic technique that solves many problems involving trees. We’ll be learning this technique by example. We’ll take a problem solving approach in this tutorial, not just describing what the final solution looks like, but walking through how one might go about solving such problems.

A first example

Task

Consider a tree in which every node has a positive value written on it. The task is to select several nodes from the tree (you may also select none), so that the sum of values of the nodes you select is maximized, and no two adjacent nodes are selected. The desired solution has complexity O(n).

General approach

As with general dynamic programming, we’re required to think of 3 things.

1. The DP state information you will keep.
2. The solution to the base case.
3. How to go from a smaller solution to a larger solution.

It is often the case that when you think of point 3, you get a answer to both points 1 and 2. But that is a vague statement, so we’ll look at a concrete example of a tree and try to calculate the answer for it.

This is the tree:

A common idea in DP on Trees is to have the sub-problems represent answers for subtrees (+ some state, we’ll see what this means in a bit). What can be a base case? Surely, it is the smallest subtree, a node. Let us think about what can be the answer for a node.

An initial solution

Consider node d. If we consider d to be a tree in its own right, then the answer for this tree is 1, as you can take d. Similarly for e and f the answers are 4 and 3 respectively.

But now think about when you go to the next bigger subtree, say subtree c. What can be the answer? If you take c, you cannot take e or f. You can take them both if you don’t take c.

Since I’m doing this exercise to illustrate a point, let us go to node a and look at possibilities. For node a, you have two possibilities, either you take it or you don’t.

• If you do, you cannot take b and c. But you can take d, e, f.
• If you do not take a, you can take b and c, and depending on you choice for them, you can / cannot take d, e, f.

Note that the logic gets into more and more cases as you go up the tree. Indeed, what we’re doing is enumerating all possible combinations of nodes we can take and take the best among them.

Where this goes wrong

But this is fundamentally flawed. The number of combinations rise exponentially with tree height. To see this just consider a “line-tree” of height h.

But there is also another point in here: Note that all logic for a subtree is derived from whether its parent node is taken / not. None of the other ancestors play a part in the calculation.

Revised solution

Inspired by this let us define our DP “state” to be not just a subtree, but a subtree + a boolean representing whether its parent is taken or not.

Now lets see what happens (we write dp[x][true/false] for indicating answer for node x, and true/false whether its parent is taken or not). Let’s start with node d again

• dp[d][false] = 1 as you can take d (parent not taken).
• dp[d][true] = 0 as you cannot take d.

The answers for e and f are also computed similarly and filled in this incomplete “DP tree:

Now lets think about c.

• dp[c][true]: c’s parent a is taken. So you cannot take c. But since you didn’t take c, e and f can be called with “parent-taken = false”. dp[c][true] = dp[e][false] + dp[f][false] = 7
• dp[c][false]: Now you have a choice.
• You can take c, then you call e and f with “parent-taken = true”, the answer in this case is 5 + dp[e][true] + dp[f][true].
• If you don’t take c, then the answer is dp[e][false] + dp[f][false] just like above.
• Since you want the best possible answer, dp[c][false] = max( 5 + dp[e][true] + dp[f][true], dp[e][false] + dp[f][false]) = max(5, 7) = 7.

The logic for b and a are similar and are filled in this “answer-tree”:

We have computed the DP and everything, but what is the final answer. Note that a can always be taken / not taken, so it effectively has a parent that is never taken. dp[a][false] = 10 is the answer.

So, in conclusion, this is our DP:

1. dp[leaf][ptaken] = v<sub>leaf</sub> if ptaken else 0
2. dp[node][true] = sum of dp[c][false] over all children c of node
3. dp[node][false] = max(v<sub>node</sub> + sum over dp[c][true] over children c of node, sum over dp[c][false] over children c of node)

Complexity of this algorithm

The time taken is proportional to sum of (number of children of v) over all nodes v of tree. Note that in this sum, each node is counted exactly once (when v = parent of this node). So the complexity is O(n = number of nodes in tree). This is way better than any exponential solution.

Simple Extension

Now let us think about how this problem could be made harder. A simple extension would be to disallow not only selection of parents, but also parent of parent.

Consider a solution of such a problem. Your DP state must carry information about parents and parents of parents. “Information” in this case would mean whether the node is taken or not. Also note that such information is sufficient, as if you know both of them, nothing else in the tree outside of this sub-tree can influence the answer for the sub-tree. In that case the state would be like dp[v][parent taken][parent-of-parent taken]. The complexity would be proportional to the size of the dp, that is 2*2*n or O(n).

In general suppose no parent within a distance of k can be taken. Then the state would be dp[parent taken][parent-of-parent taken]....[parent-of-parent-of-parent-of.... (k times) taken]. Such an algorithm will be of order O(n*2^k).

But this can be improved. Consider the state dp[v][height of closest parent taken / 0 if no parent is taken]. This takes time O(n*k).

Better Extension

Another extension can be that no two nodes at a distance of at-most 2 can be taken. This is significantly more interesting, because of this:

This cannot be captured by any state that looks like dp[v][any-information-about-parents]. Take a few minutes to see if you can come up with a solution.

Lets consider a part of a tree as shown:

Let our DP state be dp[v][1 = parent taken, 2 = parent of parent taken, 0 = none of them taken]. Suppose we’ve calculated all of them them for b, c, d, e. We want to calculate dp[a][0], dp[a][1], dp[a][2] .

• If we’re calculating dp[a][1], we cannot take a, b, c, d, e. So our answer for this case is: dp[a][1] = dp[b][2] + dp[c][2] + dp[d][2] + dp[e][2]
• If we’re calculating dp[a][2], a’s parent of parent is taken. So we can take b, c, d, e, but cannot take a. But note that we can take exactly one of b, c, d, e (Why?). So take the best among them, specifically something like, dp[a][2] = max(dp[b][0] + dp[c][2] + dp[d][2] + dp[e][2], ... (symmetrically for c, d, e, in place of b).

There is an interesting point here. Why did we take dp[c][2]? c’s parent of parent is not taken. Well if b is taken, its effect on the calculation of c is the same as if c had its parent of parent taken.

Implementation Details

Let us now think about a few implementation details:

1. DP of trees is generally implemented using a recursive depth first search (DFS). But sometimes you may also need to explicitly use a stack, if n is high.
2. States are generally stored in a multidimensional array, which is the reason why we adopted such a convention for writing answers.

In general code looks like:

dfs (v) {    
    for children c calculate dfs(c)
    calculate dp for v using children information
}    

Here is the complete code for the first problem:

#include <bits/stdc++.h>
const int maxn = 1e5 + 100; // n <= 100000 
int v[maxn]; // ensure v[i] <= 1e4, or else int won't work for dp
int dp[maxn][2]; // dynamic programming answers
std::vector<int> G[maxn]; // graph for connections
void dfs(int x, int parent) {
    for (int y: G[x]) {
        if (y == parent) continue; 
        dfs(y, x); // dfs all children
    }
    int taking_x = v[x], not_taking_x = 0; // two cases for the values
    for (int y: G[x]) {
        if (y == parent) continue;
        taking_x += dp[y][true]; 
        not_taking_x += dp[y][false]; 
    }
    dp[x][true] = not_taking_x;  // if parent of x is taken, x is not taken
    dp[x][false] = std::max(taking_x, not_taking_x);  //  if parent of x is not taken
}
int main() {
    // speed up cin/cout
    std::ios::sync_with_stdio(false); std::cin.tie(0);
    /* input start */
    int n; std::cin >> n;
    for (int i = 1; i <= n; i++) {
        std::cin >> v[i];
    }
    for (int i = 1; i < n; i++) {
        int x,y; std::cin >> x >> y;
        G[x].push_back(y); G[y].push_back(x); 
    }
  /* input ends */
    dfs(1, -1); 
    std::cout << dp[1][false] << std::endl; // 1 never has its "parent" selected.
}

Sample input:

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

Sample output:

10

You can also try out providing other inputs and check the output produced here: Code.

Practice Problems

A few practice problems on this topic:

• Problem PT07X (SPOJ)
• Problem VOCV (SPOJ)

So this is it about Dynamic Programming on Trees. Hope you enjoyed the article. Thank you for your time.