DP on Trees: Maximum Independent Set, Subtree Knapsack (O(N^2) Amortized Proof) & All-Roots Rerooting

1. Bottom-Up Subtree DP & The $O(N^2)$ Tree Knapsack Theorem

By rooting an undirected tree arbitrarily at vertex $0$, every vertex $u$ defines a subproblem on its subtree $T_u$, solved in DFS post-order:

• Maximum Weight Independent Set on a Tree (House Robber III): Maintain two states per node: $dp[u][0]$ (max weight in $T_u$ if $u$ is excluded) and $dp[u][1]$ (if $u$ is included): $$dp[u][0] = \sum_{v \in \text{children}(u)} \max(dp[v][0], dp[v][1]), \qquad dp[u][1] = w_u + \sum_{v \in \text{children}(u)} dp[v][0]$$
• Subtree Dependency Knapsack (Selecting $K$ Connected Nodes): Suppose at node $u$ we merge a knapsack table of size $\text{sz}[u]$ with child $v$'s table of size $\text{sz}[v]$. A naive bound where both loops run up to $N$ gives $O(N^3)$. However, if we bound the loop ranges strictly by the current accumulated subtree size $\text{sz}[u]$ and child size $\text{sz}[v]$ before updating $\text{sz}[u] \mathrel{+}= \text{sz}[v]$, what is the true complexity?

2. The All-Roots Rerooting Technique (Up-Down / In-Out DP)

What if we need the tree DP answer when every vertex $u \in \{0..N-1\}$ is chosen as the root (e.g., Sum of Distances in Tree, Minimum Height Trees, or Maximum Path from Every Node)? Running $N$ separate DFS traversals takes $O(N^2)$. Instead, Two-Pass Rerooting solves all $N$ roots in $O(N)$:

1. Pass 1 (Post-Order Downward DP down[u]): Root at $0$ and compute down[u] for every subtree $T_u$.
2. Pass 2 (Pre-Order Upward DP up[v]): As we step from parent $u$ to child $v$, the "complement subtree" $T \setminus T_v$ seen by $v$ consists of $u$'s upward contribution up[u] merged with all siblings $w \in \text{children}(u) \setminus \{v\}$. To avoid division (which fails on $\max/\min$ or $\bmod p$ zeroes), precompute prefix and suffix merges of $u$'s children!

3. Complexity Matrix

4. Reference C++20 Implementation

#include <algorithm>
#include <cstdint>
#include <vector>

// 1. All-Roots Rerooting DP: Compute sum of distances from EVERY node u to all other nodes in O(N)
std::vector<long long> sum_of_distances_in_tree(int n, const std::vector<std::pair<int, int>>& edges) {
    std::vector<std::vector<int>> adj(n);
    for (auto [u, v] : edges) {
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    std::vector<int> sz(n, 1);
    std::vector<long long> ans(n, 0);

    // Pass 1 (Bottom-Up): Compute subtree sizes and ans[0] = sum of depths from root 0
    auto dfs_down = [&](auto& self, int u, int p) -> void {
        for (int v : adj[u]) {
            if (v == p) continue;
            self(self, v, u);
            sz[u] += sz[v];
            ans[0] += ans[v] + sz[v]; // momentarily use ans[u] to accumulate subtree dist sum
        }
        if (u != 0) {
            long long sub_sum = 0;
            for (int v : adj[u]) if (v != p) sub_sum += ans[v] + sz[v];
            ans[u] = sub_sum;
        }
    };
    dfs_down(dfs_down, 0, -1);

    // Pass 2 (Top-Down Rerooting): Shift root from u to child v in O(1):
    // Nodes in T_v (sz[v] of them) get 1 step closer; nodes outside T_v (n - sz[v]) get 1 step farther!
    auto dfs_up = [&](auto& self, int u, int p) -> void {
        for (int v : adj[u]) {
            if (v == p) continue;
            ans[v] = ans[u] - sz[v] + (n - sz[v]);
            self(self, v, u);
        }
    };
    dfs_up(dfs_up, 0, -1);
    return ans;
}

// 2. O(N * K) Tree Dependency Knapsack (Must pick parent u to pick any node in T_u)
std::vector<long long> tree_knapsack_Exact_ONK(
    int n, int k, const std::vector<int>& parent, const std::vector<long long>& profit) {
    std::vector<std::vector<int>> children(n);
    int root = 0;
    for (int i = 0; i < n; ++i) {
        if (parent[i] == -1) root = i;
        else children[parent[i]].push_back(i);
    }

    constexpr long long NEG_INF = -(1LL << 60);
    std::vector<std::vector<long long>> dp(n);
    std::vector<int> sz(n, 0);

    auto dfs = [&](auto& self, int u) -> void {
        sz[u] = 1;
        dp[u] = { 0LL, profit[u] }; // Size 0 -> 0, Size 1 (just u) -> profit[u]

        for (int v : children[u]) {
            self(self, v);
            int new_cap = std::min(k, sz[u] + sz[v]);
            std::vector<long long> next_dp(new_cap + 1, NEG_INF);
            next_dp[0] = 0;

            // CRITICAL: Bound i by current sz[u] and j by sz[v] BEFORE adding sz[v] to sz[u]!
            for (int i = 1; i <= std::min(k, sz[u]); ++i) {
                for (int j = 0; j <= std::min(k - i, sz[v]); ++j) {
                    next_dp[i + j] = std::max(next_dp[i + j], dp[u][i] + dp[v][j]);
                }
            }
            dp[u] = std::move(next_dp);
            sz[u] += sz[v]; // Update sz[u] AFTER the merge!
        }
    };
    dfs(dfs, root);
    return dp[root];
}