Skip to content

Tree DP

Tree DP, i.e., DP performed on trees. Due to the inherent recursive nature of trees, Tree DP is generally carried out recursively.

Basics

Let's take the following problem as an example to introduce the general process of Tree DP.

Example Luogu P1352 No Boss Party

A university has \(n\) staff members, numbered \(1 \sim N\). They have subordinate relationships, meaning their relationships are like a tree rooted at the president, where a parent node is the direct supervisor of its child nodes. Now there is an anniversary celebration party. Inviting each staff member increases the happiness index by \(a_i\). However, if a staff member's direct supervisor attends the party, that staff member will definitely not attend. So, please write a program to calculate which staff members to invite to maximize the happiness index, and output the maximum happiness index.

Let \(f(i,0/1)\) represent the optimal solution for the subtree rooted at \(i\) (the second dimension value 0 means \(i\) does not attend the party, 1 means \(i\) attends the party).

For each state, there are two decisions (where \(x\) below are all children of \(i\)):

  • When the supervisor does not attend the party, subordinates may or may not attend. In this case, \(f(i,0) = \sum\max \{f(x,1),f(x,0)\}\);
  • When the supervisor attends the party, subordinates will not attend. In this case, \(f(i,1) = \sum{f(x,0)} + a_i\).

We can perform DFS and update the optimal solution for the current node when returning to the previous layer.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
#include <algorithm>
#include <iostream>
using namespace std;

struct edge {
  int v, next;
} e[6005];

int head[6005], n, cnt, f[6005][2], ans, is_h[6005], vis[6005];

void addedge(int u, int v) {  // Build graph
  e[++cnt].v = v;
  e[cnt].next = head[u];
  head[u] = cnt;
}

void calc(int k) {
  vis[k] = 1;
  for (int i = head[k]; i; i = e[i].next) {  // Enumerate each child node of the current node
    if (vis[e[i].v]) continue;
    calc(e[i].v);
    f[k][1] += f[e[i].v][0];
    f[k][0] += max(f[e[i].v][0], f[e[i].v][1]);  // State transition equation
  }
  return;
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  for (int i = 1; i <= n; i++) cin >> f[i][1];
  for (int i = 1; i < n; i++) {
    int l, k;
    cin >> l >> k;
    is_h[l] = 1;
    addedge(k, l);
  }
  for (int i = 1; i <= n; i++)
    if (!is_h[i]) {  // Start DFS from the root node
      calc(i);
      cout << max(f[i][1], f[i][0]);
      return 0;
    }
}

Typically, the state in Tree DP is the optimal solution for the current node. First, DFS traverses all optimal solutions of the subtrees, then passes them upward to the parent node for transition. Finally, the value at the root node is the desired optimal solution.

Exercises

Tree Knapsack

The knapsack problem on trees is essentially a combination of the knapsack problem and Tree DP.

Example Luogu P2014 CTSC1997 Course Selection

There are now \(n\) courses. The \(i\)-th course has credits \(a_i\). Each course has zero or one prerequisite. A course with a prerequisite must be taken after completing its prerequisite.

A student wants to take \(m\) courses. Find the maximum total credits they can obtain.

\(n,m \leq 300\)

The feature that each course has at most one prerequisite is similar to the feature in a rooted tree that a node has at most one parent.

Therefore, we can think of building a tree based on this property, so all courses form a forest structure. For convenience, we can add a course with \(0\) credits (let this course be numbered \(0\)) as the prerequisite for all courses without prerequisites. This turns the forest into a tree rooted at course \(0\).

Let \(f(u,i,j)\) represent the maximum credits obtained in the subtree rooted at \(u\), having traversed the first \(i\) subtrees of \(u\), and selected \(j\) courses.

The transition process combines the characteristics of Tree DP and Knapsack DP. We enumerate each child node \(v\) of \(u\), and simultaneously enumerate how many courses are selected in the subtree rooted at \(v\), merging the subtree results into \(u\).

Let \(s_x\) be the number of children of node \(x\), and \(\textit{siz}_x\) be the size of the subtree rooted at \(x\). The state transition equation can be written as:

\[ f(u,i,j)=\max_{v,k \leq j,k \leq \textit{siz}_v} f(u,i-1,j-k)+f(v,s_v,k) \]

Note the constraints in the state transition equation above. These constraints ensure that some meaningless states are not accessed.

The second dimension of \(f\) can be easily omitted using a rolling array. Note that in this case, \(j\) must be enumerated in reverse order.

It can be proven that the time complexity of this approach is \(O(nm)\)1.

Reference Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;
int f[305][305], s[305], n, m;
vector<int> e[305];

int dfs(int u) {
  int p = 1;
  f[u][1] = s[u];
  for (auto v : e[u]) {
    int siz = dfs(v);
    // Note the upper and lower bounds of the following two loops
    // Only consider merged subtrees, and states where the number of selected courses exceeds m+1 are meaningless
    for (int i = min(p, m + 1); i; i--)
      for (int j = 1; j <= siz && i + j <= m + 1; j++)
        f[u][i + j] = max(f[u][i + j], f[u][i] + f[v][j]);  // Transition equation
    p += siz;
  }
  return p;
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n >> m;
  for (int i = 1; i <= n; i++) {
    int k;
    cin >> k >> s[i];
    e[k].push_back(i);
  }
  dfs(0);
  cout << f[0][m + 1];
  return 0;
}

Exercises

Rerooting DP

Rerooting DP in Tree DP is also known as "second scanning". Usually, the root node is not specified, and changing the root affects some values, such as the sum of subtree depths, node weights, etc.

Typically, two DFS passes are required. The first DFS preprocesses information such as depth and node weight sums. The second DFS runs the rerooting dynamic programming.

Let's get familiar with this concept through some examples.

Example [POI2008]STA-Station

Given a tree with \(n\) nodes, find a node such that when it is the root, the sum of depths of all nodes is maximized.

Let \(u\) be the current node and \(v\) be a child node of \(u\). First, use \(s_i\) to represent the number of nodes in the subtree rooted at \(i\), with \(s_u=1+\sum s_v\). Obviously, one DFS is needed to compute all \(s_i\). This DFS is preprocessing, giving us the total number of nodes in the subtree for any given root.

Consider state transition, which is where "rerooting" comes into play. Let \(f_u\) be the sum of depths of all nodes when \(u\) is the root.

\(f_v\leftarrow f_u\) represents rerooting, i.e., transitioning from \(u\) as root to \(v\) as root. Clearly, during the rerooting transition, using \(v\) or \(u\) as the root causes changes in the depths of nodes in their subtrees. Specifically:

  • The depth of all nodes in \(v\)'s subtree decreases by one, so the total depth sum decreases by \(s_v\);

  • The depth of all nodes not in \(v\)'s subtree increases by one, so the total depth sum increases by \(n-s_v\);

Based on these two conditions, the state transition equation can be derived: \(f_v = f_u - s_v + n - s_v=f_u + n - 2 \times s_v\).

Thus, during the second DFS traversal of the tree, apply the state transition \(f_v=f_u + n - 2 \times s_v\). Then we can compute the depth sum for each node as the root. Finally, iterate through all root depth sums to find the answer.

Reference Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#include <iostream>
using namespace std;

int head[1000010 << 1], tot;
long long n, sz[1000010], dep[1000010];
long long f[1000010];

struct node {
  int to, next;
} e[1000010 << 1];

void add(int u, int v) {  // Build graph
  e[++tot] = {v, head[u]};
  head[u] = tot;
}

void dfs(int u, int fa) {  // Preprocessing dfs
  sz[u] = 1;
  dep[u] = dep[fa] + 1;
  for (int i = head[u]; i; i = e[i].next) {
    int v = e[i].to;
    if (v != fa) {
      dfs(v, u);
      sz[u] += sz[v];
    }
  }
}

void get_ans(int u, int fa) {  // Second dfs for rerooting dp
  for (int i = head[u]; i; i = e[i].next) {
    int v = e[i].to;
    if (v != fa) {
      f[v] = f[u] - sz[v] * 2 + n;
      get_ans(v, u);
    }
  }
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  int u, v;
  for (int i = 1; i <= n - 1; i++) {
    cin >> u >> v;
    add(u, v);
    add(v, u);
  }
  dfs(1, 1);
  for (int i = 1; i <= n; i++) f[1] += dep[i];
  get_ans(1, 1);
  long long int ans = -1;
  int id;
  for (int i = 1; i <= n; i++) {  // Calculate answer
    if (f[i] > ans) {
      ans = f[i];
      id = i;
    }
  }
  cout << id << '\n';
  return 0;
}

Exercises

References and Notes