Tree DP & Bitmask DP

Tree Dynamic Programming and Bitmask Dynamic Programming are advanced techniques in competitive programming. Tree DP leverages the recursive structure of trees for state transitions, while Bitmask DP uses binary representation for state sets, both solving complex combinatorial optimization problems.

Tree DP Fundamentals

The core idea of Tree DP is to use information from subtrees during post-order traversal to compute the state of the current node.

Tree Diameter - Classic Tree DP

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

class TreeDiameter {
private:
    vector<vector<pair<int, int>>> adj;
    int maxDiameter = 0;
    
public:
    TreeDiameter(int n) : adj(n) {}
    
    void addEdge(int u, int v, int w) {
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }
    
    // Returns the longest path starting from node in its subtree
    int dfs(int node, int parent) {
        int maxPath1 = 0, maxPath2 = 0;
        
        for (auto& edge : adj[node]) {
            int child = edge.first;
            int weight = edge.second;
            
            if (child == parent) continue;
            
            int childPath = dfs(child, node) + weight;
            
            // Maintain longest and second longest paths
            if (childPath > maxPath1) {
                maxPath2 = maxPath1;
                maxPath1 = childPath;
            } else if (childPath > maxPath2) {
                maxPath2 = childPath;
            }
        }
        
        // Update diameter (longest path passing through current node)
        maxDiameter = max(maxDiameter, maxPath1 + maxPath2);
        
        return maxPath1;
    }
    
    int getDiameter() {
        maxDiameter = 0;
        dfs(0, -1);
        return maxDiameter;
    }
};

// Tree Centroid
class TreeCentroid {
private:
    vector<vector<int>> adj;
    vector<int> subtreeSize;
    int n, centroid, minMaxSubtree;
    
public:
    TreeCentroid(int n) : n(n), adj(n), subtreeSize(n) {}
    
    void addEdge(int u, int v) {
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    int dfs(int node, int parent) {
        subtreeSize[node] = 1;
        int maxSubtree = 0;
        
        for (int child : adj[node]) {
            if (child == parent) continue;
            
            int childSubtreeSize = dfs(child, node);
            subtreeSize[node] += childSubtreeSize;
            maxSubtree = max(maxSubtree, childSubtreeSize);
        }
        
        // Consider the subtree above when removing current node
        maxSubtree = max(maxSubtree, n - subtreeSize[node]);
        
        if (maxSubtree < minMaxSubtree) {
            minMaxSubtree = maxSubtree;
            centroid = node;
        }
        
        return subtreeSize[node];
    }
    
    int findCentroid() {
        minMaxSubtree = n;
        dfs(0, -1);
        return centroid;
    }
};

int main() {
    // Test tree diameter
    TreeDiameter td(6);
    td.addEdge(0, 1, 2);
    td.addEdge(0, 2, 3);
    td.addEdge(1, 3, 1);
    td.addEdge(1, 4, 4);
    td.addEdge(2, 5, 2);
    
    cout << "Tree diameter: " << td.getDiameter() << endl;
    
    // Test tree centroid
    TreeCentroid tc(6);
    tc.addEdge(0, 1);
    tc.addEdge(0, 2);
    tc.addEdge(1, 3);
    tc.addEdge(1, 4);
    tc.addEdge(2, 5);
    
    cout << "Tree centroid: " << tc.findCentroid() << endl;
    
    return 0;
}

Tree diameter uses Tree DP to maintain longest and second longest paths at each node. Tree centroid is the node whose removal results in the smallest maximum connected component.

Rerooting DP Technique

Rerooting DP first computes answers with any node as root, then uses state transitions to compute answers for other nodes as roots.

Rerooting DP - Sum of Distances from Each Node

#include <iostream>
#include <vector>
using namespace std;

class RerootingDP {
private:
    vector<vector<int>> adj;
    vector<long long> subtreeSum;    // Sum of distances in subtree
    vector<int> subtreeSize;         // Size of subtree
    vector<long long> ans;           // Answer for each node
    int n;
    
public:
    RerootingDP(int n) : n(n), adj(n), subtreeSum(n), subtreeSize(n), ans(n) {}
    
    void addEdge(int u, int v) {
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    // First DFS: compute subtree information with node as root
    void dfs1(int node, int parent) {
        subtreeSize[node] = 1;
        subtreeSum[node] = 0;
        
        for (int child : adj[node]) {
            if (child == parent) continue;
            
            dfs1(child, node);
            subtreeSize[node] += subtreeSize[child];
            subtreeSum[node] += subtreeSum[child] + subtreeSize[child];
        }
    }
    
    // Second DFS: reroot to compute answer for each node
    void dfs2(int node, int parent) {
        for (int child : adj[node]) {
            if (child == parent) continue;
            
            // Compute answer when child is root
            // Transfer from node to child:
            // 1. Child's subtree contribution remains same
            // 2. Rest (including node) contribution needs recalculation
            
            long long childSubtreeSum = subtreeSum[child];
            int childSubtreeSize = subtreeSize[child];
            
            // Remove child's contribution from node
            subtreeSum[node] -= (childSubtreeSum + childSubtreeSize);
            subtreeSize[node] -= childSubtreeSize;
            
            // Add node's contribution to child
            subtreeSum[child] += (subtreeSum[node] + subtreeSize[node]);
            subtreeSize[child] += subtreeSize[node];
            
            ans[child] = subtreeSum[child];
            
            dfs2(child, node);
            
            // Restore original values
            subtreeSum[child] = childSubtreeSum;
            subtreeSize[child] = childSubtreeSize;
            subtreeSum[node] += (childSubtreeSum + childSubtreeSize);
            subtreeSize[node] += childSubtreeSize;
        }
    }
    
    vector<long long> solve() {
        dfs1(0, -1);
        ans[0] = subtreeSum[0];
        dfs2(0, -1);
        return ans;
    }
};

int main() {
    RerootingDP rd(6);
    rd.addEdge(0, 1);
    rd.addEdge(0, 2);
    rd.addEdge(1, 3);
    rd.addEdge(1, 4);
    rd.addEdge(2, 5);
    
    vector<long long> result = rd.solve();
    
    cout << "Sum of distances from each node:" << endl;
    for (int i = 0; i < result.size(); i++) {
        cout << "Node " << i << ": " << result[i] << endl;
    }
    
    return 0;
}

Rerooting DP avoids recomputing from scratch for each root. It transfers states efficiently by considering how contributions change when moving the root.

Bitmask DP Fundamentals

Bitmask DP uses binary representation to encode state sets, suitable for problems with exponential state space but small parameter ranges.

Traveling Salesman Problem - Classic Bitmask DP

#include <iostream>
#include <vector>
#include <algorithm>
#include <climits>
using namespace std;

class TSP {
private:
    vector<vector<int>> dist;
    vector<vector<int>> dp;
    int n;
    
public:
    TSP(vector<vector<int>>& distances) : dist(distances) {
        n = dist.size();
        dp.assign(1 << n, vector<int>(n, INT_MAX));
    }
    
    int solve() {
        // Base case: start from city 0
        dp[1][0] = 0; // mask = 1 (only city 0 visited), current city = 0
        
        for (int mask = 1; mask < (1 << n); mask++) {
            for (int u = 0; u < n; u++) {
                if (!(mask & (1 << u)) || dp[mask][u] == INT_MAX) continue;
                
                for (int v = 0; v < n; v++) {
                    if (mask & (1 << v)) continue; // City v already visited
                    
                    int newMask = mask | (1 << v);
                    dp[newMask][v] = min(dp[newMask][v], dp[mask][u] + dist[u][v]);
                }
            }
        }
        
        // Find minimum cost to return to starting city
        int result = INT_MAX;
        int finalMask = (1 << n) - 1; // All cities visited
        
        for (int u = 1; u < n; u++) {
            if (dp[finalMask][u] != INT_MAX) {
                result = min(result, dp[finalMask][u] + dist[u][0]);
            }
        }
        
        return result == INT_MAX ? -1 : result;
    }
    
    vector<int> getPath() {
        vector<int> path;
        int mask = (1 << n) - 1;
        int current = 0;
        
        // Find the ending city that gives minimum cost
        for (int u = 1; u < n; u++) {
            if (dp[mask][u] + dist[u][0] < dp[mask][current] + dist[current][0]) {
                current = u;
            }
        }
        
        // Reconstruct path
        while (mask != 1) {
            path.push_back(current);
            int prevMask = mask ^ (1 << current);
            
            for (int u = 0; u < n; u++) {
                if ((prevMask & (1 << u)) && 
                    dp[prevMask][u] + dist[u][current] == dp[mask][current]) {
                    mask = prevMask;
                    current = u;
                    break;
                }
            }
        }
        
        path.push_back(0);
        reverse(path.begin(), path.end());
        return path;
    }
};

// Assignment Problem using Bitmask DP
class AssignmentProblem {
private:
    vector<vector<int>> cost;
    vector<int> dp;
    int n;
    
public:
    AssignmentProblem(vector<vector<int>>& costs) : cost(costs) {
        n = cost.size();
        dp.assign(1 << n, INT_MAX);
    }
    
    int solve() {
        dp[0] = 0; // No tasks assigned
        
        for (int mask = 0; mask < (1 << n); mask++) {
            if (dp[mask] == INT_MAX) continue;
            
            int person = __builtin_popcount(mask); // Number of assigned tasks
            if (person >= n) continue;
            
            for (int task = 0; task < n; task++) {
                if (mask & (1 << task)) continue; // Task already assigned
                
                int newMask = mask | (1 << task);
                dp[newMask] = min(dp[newMask], dp[mask] + cost[person][task]);
            }
        }
        
        return dp[(1 << n) - 1];
    }
};

int main() {
    // Test TSP
    vector<vector<int>> distances = {
        {0, 10, 15, 20},
        {10, 0, 35, 25},
        {15, 35, 0, 30},
        {20, 25, 30, 0}
    };
    
    TSP tsp(distances);
    int minCost = tsp.solve();
    vector<int> path = tsp.getPath();
    
    cout << "TSP minimum cost: " << minCost << endl;
    cout << "Optimal path: ";
    for (int i = 0; i < path.size(); i++) {
        cout << path[i];
        if (i < path.size() - 1) cout << " -> ";
    }
    cout << " -> 0" << endl;
    
    // Test Assignment Problem
    vector<vector<int>> costs = {
        {9, 2, 7, 8},
        {6, 4, 3, 7},
        {5, 8, 1, 8},
        {7, 6, 9, 4}
    };
    
    AssignmentProblem ap(costs);
    cout << "\nAssignment minimum cost: " << ap.solve() << endl;
    
    return 0;
}

Bitmask DP encodes visited states using binary bits. TSP uses dp[mask][u] representing minimum cost to visit cities in mask ending at u. Assignment problem assigns tasks to people optimally.

Advanced Tree DP Applications

Advanced Tree DP handles complex constraints and multiple state dimensions, often combined with other techniques.

Tree Coloring with Constraints

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

class TreeColoring {
private:
    vector<vector<int>> adj;
    vector<vector<long long>> dp;
    int n, k;
    const int MOD = 1000000007;
    
public:
    TreeColoring(int n, int k) : n(n), k(k), adj(n), dp(n, vector<long long>(k, 0)) {}
    
    void addEdge(int u, int v) {
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    // dp[u][c] = number of ways to color subtree rooted at u with u colored c
    void dfs(int u, int parent) {
        // Base case: leaf node can be colored in any color
        for (int c = 0; c < k; c++) {
            dp[u][c] = 1;
        }
        
        for (int v : adj[u]) {
            if (v == parent) continue;
            
            dfs(v, u);
            
            // Combine results from child v
            vector<long long> newDp(k, 0);
            
            for (int uc = 0; uc < k; uc++) {
                for (int vc = 0; vc < k; vc++) {
                    if (uc != vc) { // Adjacent nodes must have different colors
                        newDp[uc] = (newDp[uc] + dp[u][uc] * dp[v][vc]) % MOD;
                    }
                }
            }
            
            dp[u] = newDp;
        }
    }
    
    long long solve() {
        dfs(0, -1);
        long long result = 0;
        for (int c = 0; c < k; c++) {
            result = (result + dp[0][c]) % MOD;
        }
        return result;
    }
};

// Maximum Weight Independent Set in Tree
class MaxWeightIndependentSet {
private:
    vector<vector<int>> adj;
    vector<int> weight;
    vector<vector<long long>> dp;
    int n;
    
public:
    MaxWeightIndependentSet(vector<int>& weights) : weight(weights) {
        n = weight.size();
        adj.resize(n);
        dp.assign(n, vector<long long>(2, 0));
    }
    
    void addEdge(int u, int v) {
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    // dp[u][0] = max weight when u is not selected
    // dp[u][1] = max weight when u is selected
    void dfs(int u, int parent) {
        dp[u][0] = 0;
        dp[u][1] = weight[u];
        
        for (int v : adj[u]) {
            if (v == parent) continue;
            
            dfs(v, u);
            
            // If u is not selected, v can be either selected or not
            dp[u][0] += max(dp[v][0], dp[v][1]);
            
            // If u is selected, v cannot be selected
            dp[u][1] += dp[v][0];
        }
    }
    
    long long solve() {
        dfs(0, -1);
        return max(dp[0][0], dp[0][1]);
    }
    
    vector<int> getSelectedNodes() {
        vector<int> selected;
        function<void(int, int, bool)> reconstruct = [&](int u, int parent, bool canSelect) {
            if (canSelect && dp[u][1] >= dp[u][0]) {
                selected.push_back(u);
                for (int v : adj[u]) {
                    if (v != parent) {
                        reconstruct(v, u, false);
                    }
                }
            } else {
                for (int v : adj[u]) {
                    if (v != parent) {
                        reconstruct(v, u, true);
                    }
                }
            }
        };
        
        reconstruct(0, -1, true);
        return selected;
    }
};

int main() {
    // Test tree coloring
    TreeColoring tc(5, 3);
    tc.addEdge(0, 1);
    tc.addEdge(0, 2);
    tc.addEdge(1, 3);
    tc.addEdge(1, 4);
    
    cout << "Tree coloring ways with 3 colors: " << tc.solve() << endl;
    
    // Test maximum weight independent set
    vector<int> weights = {10, 5, 2, 7, 8};
    MaxWeightIndependentSet mwis(weights);
    mwis.addEdge(0, 1);
    mwis.addEdge(0, 2);
    mwis.addEdge(1, 3);
    mwis.addEdge(1, 4);
    
    cout << "Maximum weight independent set: " << mwis.solve() << endl;
    
    vector<int> selected = mwis.getSelectedNodes();
    cout << "Selected nodes: ";
    for (int node : selected) {
        cout << node << " ";
    }
    cout << endl;
    
    return 0;
}

Advanced Tree DP handles multiple constraints simultaneously. Tree coloring ensures adjacent nodes have different colors. Maximum weight independent set selects non-adjacent nodes with maximum total weight.

Competitive Programming Techniques

Tree DP Core Concepts:

Utilize tree's recursive structure for bottom-up state computation
Define states based on subtree properties and node characteristics
Combine child states to compute parent state during post-order traversal
Consider rerooting for problems requiring all nodes as potential roots

Bitmask DP Applications:

Exponential state space with small parameter constraints (n ≤ 20)
Subset enumeration and selection problems
State compression using binary representation
Efficient transitions using bitwise operations

Implementation Tips:

Tree DP: Handle parent-child relationships correctly during DFS
Bitmask DP: Use __builtin_popcount() for bit counting
Memory optimization: Use rolling arrays when possible
State reconstruction: Track transition choices for solution recovery

Common Patterns:

Tree diameter, centroid, and distance problems
TSP, assignment, and subset selection problems
Tree coloring and independent set problems
Counting problems with tree or subset constraints