DP on Trees

What is DP on Trees?

Dynamic Programming (DP) on Trees is a technique used to solve optimization problems on tree data structures. Unlike linear DP where we iterate over arrays, tree DP processes the tree in a post-order (bottom-up) or pre-order (top-down) fashion. It often involves storing results for subtrees at each node and combining them to solve larger subproblems.

General Algorithm

1. Choose a root (arbitrary for unrooted trees).
2. Traverse the tree using DFS (post-order for bottom-up DP).
3. At each node, compute the result based on its children's results.
4. Store the result in a DP array indexed by node.

Time and Memory Complexity

• Time Complexity: O(N), where N is the number of nodes. Each node and edge is visited once.
• Space Complexity: O(N) for the tree representation and the DP array.

Easy Problem 1: Find Height of Each Subtree

Problem: Given a tree, find the height of each node's subtree.

#include <iostream>
#include <vector>

using namespace std;

const int N = 100005;
vector<int> tree[N];
int height[N];

void dfs(int node, int parent) {
    height[node] = 0;
    for (int child : tree[node]) {
        if (child == parent) continue;
        dfs(child, node);
        height[node] = max(height[node], 1 + height[child]);
    }
}

int main() {
    int n;
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(1, -1);
    for (int i = 1; i <= n; i++) {
        cout << "Height of node " << i << " is " << height[i] << endl;
    }
}

Easy Problem 2: Count Number of Nodes in Subtree

#include <iostream>
#include <vector>

using namespace std;

const int N = 100005;
vector<int> tree[N];
int subtree_size[N];

void dfs(int node, int parent) {
    subtree_size[node] = 1;
    for (int child : tree[node]) {
        if (child == parent) continue;
        dfs(child, node);
        subtree_size[node] += subtree_size[child];
    }
}

int main() {
    int n;
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(1, -1);
    for (int i = 1; i <= n; i++) {
        cout << "Subtree size of node " << i << " is " << subtree_size[i] << endl;
    }
}

Intermediate Problem 1: Diameter of Tree

#include <iostream>
#include <vector>

using namespace std;

const int N = 100005;
vector<int> tree[N];
int max_diameter = 0;

int dfs(int node, int parent) {
    int max1 = 0, max2 = 0;
    for (int child : tree[node]) {
        if (child == parent) continue;
        int depth = dfs(child, node) + 1;
        if (depth > max1) {
            max2 = max1;
            max1 = depth;
        } else if (depth > max2) {
            max2 = depth;
        }
    }
    max_diameter = max(max_diameter, max1 + max2);
    return max1;
}

int main() {
    int n;
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(1, -1);
    cout << "Diameter of the tree is " << max_diameter << endl;
}

Intermediate Problem 2: Tree DP - Max Sum of Independent Set

Select nodes such that no two adjacent nodes are selected and the sum is maximum.

#include <iostream>
#include <vector>

using namespace std;

const int N = 100005;
vector<int> tree[N];
int value[N], dp[N][2];

void dfs(int node, int parent) {
    dp[node][0] = 0;
    dp[node][1] = value[node];

    for (int child : tree[node]) {
        if (child == parent) continue;
        dfs(child, node);
        dp[node][0] += max(dp[child][0], dp[child][1]);
        dp[node][1] += dp[child][0];
    }
}

int main() {
    int n;
    cin >> n;
    for (int i = 1; i <= n; i++) cin >> value[i];
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(1, -1);
    cout << "Max sum of independent set: " << max(dp[1][0], dp[1][1]) << endl;
}

Hard Problem: Rerooting Technique - Sum of Distances in Tree

#include <iostream>
#include <vector>

using namespace std;

const int N = 100005;
vector<int> tree[N];
int dp[N], subtree_size[N], ans[N];
int n;

void dfs1(int node, int parent) {
    subtree_size[node] = 1;
    for (int child : tree[node]) {
        if (child == parent) continue;
        dfs1(child, node);
        subtree_size[node] += subtree_size[child];
        dp[node] += dp[child] + subtree_size[child];
    }
}

void dfs2(int node, int parent) {
    for (int child : tree[node]) {
        if (child == parent) continue;
        ans[child] = ans[node] - subtree_size[child] + (n - subtree_size[child]);
        dfs2(child, node);
    }
}

int main() {
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs1(1, -1);
    ans[1] = dp[1];
    dfs2(1, -1);
    for (int i = 1; i <= n; i++) {
        cout << "Sum of distances from node " << i << ": " << ans[i] << endl;
    }
}

Conclusion

DP on Trees is a powerful technique to solve many hierarchical or recursive structure problems efficiently. By using bottom-up or rerooting strategies, even complex dependencies in trees can be handled within linear time.