Skip to content
Embercode
✕
DP optimisationsLesson 4 of 11

45 min read

DP on trees and DAGs

Why this matters when you build with AI

Trees and graphs appear in almost every system that stores hierarchies: org charts, file folders, dependency graphs, build pipelines and category trees. An AI tool asked to compute something over such a structure usually writes a recursive function. The recursion is short and readable, and it works on small tests. It fails on the large inputs that real data produces.

I tested this with a short separate program. A recursive walk down a straight line of 100,000 nodes finished. The same walk on 1,000,000 nodes ended with exit code 139, which is a segmentation fault. That means the program ran past the memory set aside for its call stack. On this machine the default stack size is 8192 KB. The limit depends on the machine, so your result may differ. The safe habit is to visit nodes in a flat order that you build with a queue, as the code below does.

A second common error is rerooting by running a full walk from every node. For each node that costs O(n) steps, so the total is O(n^2). With n = 100,000 that is 10^10 steps, which is too slow. Rerooting gets every answer in two passes.

A third error is a DAG program that runs without checking for cycles. A cycle is a loop of arrows, such as 1 to 2 and 2 back to 1. A DAG (directed acyclic graph) has no such loop. In the lesson's cycle test, the program printed count 0 for a node and a longest path of minus a huge number, with no error message. A check for missing nodes after the topological sort catches this, and the mistakes section gives the check.

Terms you need first

  • Tree: a set of n nodes joined by n - 1 edges, with no loop. Pick one node as the root. The parent of a node is the neighbour on the way to the root. The children are the other neighbours. The subtree of a node is the node and everything below it.
  • Subtree DP: a dynamic programming method (storing answers for states and reusing them) that computes a value for each subtree from the values of its children. It works from the leaves up, so each child is done before its parent.
  • Rerooting: computing an answer for every node, as if that node were the root. It uses one pass down the tree and one pass back up.
  • DAG: a directed graph with no cycle. Each arrow points one way.
  • Topological order: an order of the nodes in which every arrow goes from an earlier node to a later node. A DAG always has one.
  • Knapsack table: for a subtree, a list where entry j is the best value of a selection that has exactly j nodes. The name comes from the 0/1 knapsack problem, where you choose items to fit a limit.

The idea: subtree DP with a team example

Take this problem, which I wrote for this lesson. A company has n people in a reporting tree. Person 0 is the top manager, who is the root. A team is valid if it contains person 0, and whenever it contains a person, it also contains that person's manager. Each person has a skill score. For each team size j from 1 to n, find the largest total skill of a valid team with exactly j people.

Work through a small tree. Person 0 manages persons 1 and 2. Person 1 manages persons 3 and 4. Person 2 manages persons 5, 6 and 7. The skills are 4, 7, 2, 9, 3, 6, 8 and 5.

Start at the leaves, the people with no team of their own. Person 3 has skill 9. Her knapsack table is: size 1 gives 9, since she alone is a valid team if we count her as the top of her own subtree. Now person 1 (skill 7) with his children 3 and 4 (skill 3). Let us count how the table grows.

  • Start with only person 1: size 1 gives 7.
  • Add child 3. Either take nothing from her, or take her team. Taking her team of size 1 with value 9 gives size 2 with value 7 + 9 = 16.
  • Add child 4. Taking nothing keeps the old values. Taking her team gives size 3 with value 16 + 3 = 19. Size 2 stays 16, because the best size-2 team is still persons 1 and 3.

So the table for person 1 is: size 1 gives 7, size 2 gives 16, size 3 gives 19. A parent combines its own table with each child's table in the same way. The combining rule is the whole idea: for each pair of sizes a from the current table and b from the child table, the new size a + b gets the sum of the two values, and b = 0 means "take nothing from this child".

Here is that step as a formula. Let cur be the table so far, and let g be the child table with g[0] = 0. Then new[a + b] = max(new[a + b], cur[a] + g[b]) for every valid a and b. The first index a starts at 1, because the table is only for selections that contain the current node.

The code below does this with a queue-built order, so it never recurses. It returns the best value for each size j from 1 to n for the whole tree.

Standard technique and the size bound

Two pieces of code carry this lesson. The first is the knapsack merge. The second is rerooting for a weighted tree. The third is the DAG code. All three live in one program, and the program's output is pasted after it.

The program was compiled with g++ -std=c++17 -O2 -Wall in a scratch folder. The knapsack was checked against all subsets of up to 10 nodes, the eccentricity against a walk from each node, and the DAG counts and longest paths against a list of every path. These checks used random small inputs, and 4,851 checks agreed.

// Tree knapsack merge, rerooting, and DAG DP in topological order.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll NEG = LLONG_MIN / 4;
 
// BFS order and parent array from a root. Children appear after their parent in `order`.
void bfs_order(int n, const vector<vector<int>> &adj, int root, vector<int> &order, vector<int> &parent) {
    order.clear();
    parent.assign(n, -1);
    vector<char> seen(n, 0);
    order.push_back(root);
    seen[root] = 1;
    for (size_t i = 0; i < order.size(); i++) {
        int v = order[i];
        for (int u : adj[v])
            if (!seen[u]) {
                seen[u] = 1;
                parent[u] = v;
                order.push_back(u);
            }
    }
}
 
// Best value of a connected set that contains the root and has exactly j nodes, for j = 1..n.
// dp[v][j] = best value of a set inside subtree(v) that contains v and has j nodes.
// Returns res with res[j] for j = 1..n (res[0] is unused).
vector<ll> tree_knapsack(int n, const vector<vector<int>> &adj, const vector<ll> &val, int root, ll &merge_steps) {
    vector<int> order, parent;
    bfs_order(n, adj, root, order, parent);
    vector<vector<ll>> dp(n);
    merge_steps = 0;
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        vector<ll> cur = {NEG, val[v]};  // size 1: only v itself
        for (int c : adj[v]) {
            if (c == parent[v]) continue;
            // Child option g: index 0 means "take nothing from this child".
            vector<ll> &g = dp[c];
            g[0] = 0;
            vector<ll> nxt(cur.size() + g.size() - 1, NEG);
            for (size_t a = 1; a < cur.size(); a++) {
                if (cur[a] == NEG) continue;
                for (size_t b = 0; b < g.size(); b++) {
                    merge_steps++;
                    if (g[b] == NEG) continue;
                    nxt[a + b] = max(nxt[a + b], cur[a] + g[b]);
                }
            }
            cur.swap(nxt);
            vector<ll>().swap(g);  // free the child's memory
        }
        dp[v].swap(cur);
    }
    vector<ll> res(dp[root].begin(), dp[root].end());
    res.resize(n + 1, NEG);
    return res;
}
 
// Rerooting: ecc[v] = max over u of dist(v, u), on a weighted tree.
// down[v]: farthest distance from v into its own subtree (0 if none).
// up[v]: farthest distance from v to a node outside its subtree (0 at the root, which
// only stands for v itself and never changes the maximum).
vector<ll> eccentricity(int n, const vector<vector<pair<int, ll>>> &adj) {
    vector<vector<int>> plain(n);
    for (int v = 0; v < n; v++)
        for (auto [u, w] : adj[v]) plain[v].push_back(u);
    vector<int> order, parent;
    bfs_order(n, plain, 0, order, parent);
    vector<ll> pw(n, 0);  // weight of the edge to the parent
    for (int v = 0; v < n; v++)
        for (auto [u, w] : adj[v])
            if (parent[u] == v) pw[u] = w;
    vector<ll> down(n, 0), up(n, 0);
    vector<int> best_child(n, -1);
    vector<ll> best1(n, 0), best2(n, 0);  // top two (down[c] + w) values, 0 means v itself
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        down[v] = best1[v];
        if (parent[v] >= 0) {
            int p = parent[v];
            ll cand = down[v] + pw[v];
            if (cand > best1[p]) {
                best2[p] = best1[p];
                best1[p] = cand;
                best_child[p] = v;
            } else if (cand > best2[p]) {
                best2[p] = cand;
            }
        }
    }
    for (int i = 0; i < n; i++) {
        int v = order[i];
        for (auto [c, w] : adj[v]) {
            if (c == parent[v]) continue;
            ll excl = (best_child[v] == c) ? best2[v] : best1[v];
            up[c] = w + max(up[v], excl);
        }
    }
    vector<ll> ecc(n);
    for (int v = 0; v < n; v++) ecc[v] = max(down[v], up[v]);
    return ecc;
}
 
// Longest path and number of paths in a DAG (edges go from a to b), computed in topological order.
// Kahn's algorithm: repeatedly take a node with no unprocessed incoming edge.
struct DagResult {
    vector<ll> longest;  // longest[v]: largest weight of a path that starts at source and ends at v
    vector<ll> count;    // count[v]: number of source-to-v paths, mod MOD
    vector<int> from;    // predecessor on the longest path, -1 at the source
};
 
DagResult dag_paths(int n, const vector<vector<pair<int, ll>>> &adj, int source, ll MOD) {
    vector<int> indeg(n, 0);
    for (int v = 0; v < n; v++)
        for (auto [u, w] : adj[v]) indeg[u]++;
    vector<int> q;
    for (int v = 0; v < n; v++)
        if (indeg[v] == 0) q.push_back(v);
    vector<int> topo;
    for (size_t i = 0; i < q.size(); i++) {
        int v = q[i];
        topo.push_back(v);
        for (auto [u, w] : adj[v])
            if (--indeg[u] == 0) q.push_back(u);
    }
    DagResult r;
    r.longest.assign(n, NEG);
    r.count.assign(n, 0);
    r.from.assign(n, -1);
    r.longest[source] = 0;
    r.count[source] = 1;
    for (int v : topo) {
        if (r.longest[v] == NEG) continue;  // not reachable from source
        for (auto [u, w] : adj[v]) {
            if (r.longest[v] + w > r.longest[u]) {
                r.longest[u] = r.longest[v] + w;
                r.from[u] = v;
            }
            r.count[u] = (r.count[u] + r.count[v]) % MOD;
        }
    }
    return r;
}
 
int main() {
    // Demo tree: 8 nodes, root 0 is the team lead. val[v] = skill score of person v.
    int n = 8;
    vector<vector<int>> adj(n);
    auto edge = [&](int a, int b) { adj[a].push_back(b); adj[b].push_back(a); };
    edge(0, 1); edge(0, 2); edge(1, 3); edge(1, 4); edge(2, 5); edge(5, 6); edge(5, 7);
    vector<ll> val = {4, 7, 2, 9, 3, 6, 8, 5};
    ll steps = 0;
    vector<ll> best = tree_knapsack(n, adj, val, 0, steps);
    printf("knapsack, best total for a root-connected team of size j:\n");
    for (int j = 1; j <= n; j++) printf("j=%d best=%lld\n", j, best[j]);
    printf("merge steps: %lld\n", steps);
 
    // Weighted tree for rerooting: edge weights are road lengths.
    int m = 6;
    vector<vector<pair<int, ll>>> g(m);
    auto wedge = [&](int a, int b, ll w) { g[a].push_back({b, w}); g[b].push_back({a, w}); };
    wedge(0, 1, 3); wedge(1, 2, 4); wedge(1, 3, 1); wedge(3, 4, 6); wedge(0, 5, 2);
    vector<ll> ecc = eccentricity(m, g);
    printf("eccentricity:");
    for (int v = 0; v < m; v++) printf(" %lld", ecc[v]);
    printf("\n");
 
    // DAG of tasks: edge a->b with weight w means task b starts after a with cost w.
    int k = 6;
    vector<vector<pair<int, ll>>> d(k);
    auto dedge = [&](int a, int b, ll w) { d[a].push_back({b, w}); };
    dedge(0, 1, 2); dedge(0, 2, 5); dedge(1, 3, 4); dedge(2, 3, 1); dedge(3, 4, 3); dedge(1, 4, 9); dedge(4, 5, 2);
    DagResult res = dag_paths(k, d, 0, 1000000007LL);
    printf("dag longest to 5: %lld via", res.longest[5]);
    for (int v = 5; v != -1; v = res.from[v]) printf(" %d", v);
    printf("\ndag paths to 5 mod p: %lld\n", res.count[5]);
    return 0;
}

Real output of that program:

knapsack, best total for a root-connected team of size j:
j=1 best=4
j=2 best=11
j=3 best=20
j=4 best=23
j=5 best=28
j=6 best=36
j=7 best=41
j=8 best=44
merge steps: 40
eccentricity: 10 7 11 6 12 12
dag longest to 5: 13 via 5 4 1 0
dag paths to 5 mod p: 3

The knapsack result can be checked by hand. For size 3 the best team is persons 0, 1 and 3, with skills 4 + 7 + 9 = 20. Size 2 is persons 0 and 1, with 11. Size 8 is everybody, with the total of all skills, 44.

Why the merge is O(n^2) in total

Merging a table of size a with a table of size b takes about a times b steps. The bound comes from counting pairs. Two nodes of the tree are combined in exactly one merge, at their lowest common ancestor (the lowest node that is an ancestor of both). So the total number of combined pairs is at most n times (n - 1) divided by 2. The sum of all the products a times b over the whole tree is therefore O(n^2).

The program counts its inner loop steps in merge_steps. I measured it on two shapes. For a straight line of 4000 nodes it counted 8,001,999 steps, which is 0.50 times n^2. For a balanced binary tree of 4000 nodes it counted 8,022,181 steps, which is 0.50 times n^2 as well. For a star with n = 5, a root and four leaves, it counted 20 steps, which is n times (n - 1). These counts include the b = 0 option, so they are a little above the number of real pairs.

A caution on the bound. It holds when each merge loops over the sizes that the two tables can really have. A version that loops over all n sizes for every merge costs O(n) per merge, and O(n^2) merges give O(n^3). The loop bounds in the code above are the table sizes, which is what makes the total O(n^2).

Rerooting

Suppose we want, for every node v, the largest weighted distance from v to any other node. This is the eccentricity of v. A simple method runs a walk from each node, which costs O(n) per node and O(n^2) in total. Rerooting does it in O(n).

Use two numbers per node. down[v] is the farthest distance from v to a node inside its own subtree. up[v] is the farthest distance from v to a node outside its subtree. The eccentricity is the larger of the two.

The down pass runs from the leaves up. down[v] is the best value of down[c] + w over the children c, where w is the edge weight to the child. The up pass runs from the root down. For a child c of v with edge weight w, the nodes outside c's subtree are reached through v. The farthest of them is at distance w plus the best of three choices: v itself (distance 0), the farthest node through v's parent (up[v]), or the farthest node through another child of v.

To make "the farthest node through another child" fast, the code keeps the top two child values for each node. If c is the child that gave the top value, it uses the second value. Otherwise it uses the top value.

Check on the demo tree. Node 0 has eccentricity 10, from node 0 to node 4 through node 1 and node 3, which is 3 + 1 + 6 = 10. The program prints 10 for node 0.

DAG in topological order

A DAG has a topological order: a list of all nodes in which every arrow goes forward. Kahn's algorithm builds one. It counts the incoming arrows of each node. It puts every node with no incoming arrows in a queue. It takes a node off the queue, appends it to the order, and removes its outgoing arrows, which may free more nodes.

Once the order exists, one pass over the nodes computes everything. When we reach a node, all of its incoming arrows have already been handled. So its longest distance and its path count are final, and it can pass them on to its neighbours.

In the demo, the longest path from node 0 to node 5 is 13, through nodes 0, 1, 4 and 5 with weights 2, 9 and 2. The printed path lists the nodes from the end back to the start, so it reads 5, 4, 1, 0. There are 3 paths from 0 to 5: 0-1-4-5, 0-2-3-4-5 and 0-1-3-4-5.

For the path count the program uses % MOD on every addition. The counts can grow very large, so a modulus keeps them in range. The longest path is not modded, because it is a sum of weights and is compared directly.

Complexity

  • Tree knapsack: O(n^2) time in total, as shown above. Memory is O(n): the tables that are alive at one moment belong to disjoint subtrees, so their sizes add up to at most n plus the number of tables.
  • Rerooting: O(n) time and O(n) memory, two passes over the tree.
  • DAG paths: O(n + m) time, where m is the number of edges, and O(n + m) memory.

This lesson gives no Python version of the knapsack. Its O(n^2) merge runs many inner-loop steps, and Python executes each step much more slowly than C++. The C++ version is the one to use when n is in the thousands or more. The exact time depends on the machine.

Common mistakes

  1. Recursion on a long path. A recursive walk over a path of 1,000,000 nodes ended with exit code 139 (a segmentation fault) on this machine, while 100,000 nodes worked. The fix is the queue-based order in bfs_order, which the code above uses for every tree walk.

  2. Rerooting with a full walk per node. The code is correct on 10 nodes but takes about 10^10 steps at n = 100,000, so the symptom is a time limit, not a wrong answer. Use two passes instead.

  3. Forgetting the node itself in the up pass. If up[c] only looks at the parent's up and the other children, the parent itself is left out, and a node whose farthest target is the parent can get an answer that is too small. The code adds 0 for the node itself, which is why the root's up value of 0 is safe.

  4. Reading a table after it was freed. A child's table dp[c] is freed after its merge, so reading it again gives an empty vector and undefined behaviour. The fix is to merge each child exactly once, which the loop over adj[v] does.

  5. DAG code that does not check for cycles. The next example has a cycle, 1 to 2 and 2 back to 1. The output with the code above was:

node 0 longest=0 count=1
node 1 longest=2 count=1
node 2 longest=-2305843009213693952 count=0

Node 2 is on a cycle, so it is never added to the topological order. Its longest value is the NEG sentinel and its count is 0. Nothing crashes. The fix is to count the nodes in topo after Kahn's algorithm. If topo.size() < n, the graph has a cycle, and the answers must not be used.

  1. Mixing up the two indices of a table. The table index is the number of nodes in the selection, not the total value. An off-by-one here shifts every answer by one size. Check size 1, which must equal the root's own value.

Hands-on exercise

Write your own tree knapsack that counts instead of maximising. For each size j, count the valid teams of size j, that is, the connected sets that contain the root with exactly j nodes, modulo 1000000007. Change the operations in the merge from max and plus to plus and times, and start the table with cur = {0, 1} and g[0] = 1.

Test it. Write a brute force that tries every subset of up to 10 nodes, keeps the sets that contain the root and are closed under taking managers, and counts them by size. Run both on 200 random trees with n from 1 to 10, and compare every size. Use the same random generator for both, so a failure is easy to reproduce.

A correct version agrees on every size. If one size fails, print the tree and the subset that the brute force found, and compare the two by hand on that tree.

Show the answer

The changes from the maximum version are the start cur = {0, 1}, g[0] = 1 for taking nothing, and the merge line nxt[a + b] = (nxt[a + b] + cur[a] * g[b]) % MOD. This function is the one the program below uses. It was compiled with g++ -std=c++17 -O2 -Wall and run on the demo tree from this lesson.

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
 
vector<ll> count_root_sets(int n, const vector<vector<int>> &adj, ll MOD) {
    vector<int> order, par(n, -1);
    vector<char> seen(n, 0);
    order.push_back(0); seen[0] = 1;
    for (size_t i = 0; i < order.size(); i++)
        for (int u : adj[order[i]]) if (!seen[u]) { seen[u] = 1; par[u] = order[i]; order.push_back(u); }
    vector<vector<ll>> dp(n);
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        vector<ll> cur = {0, 1};
        for (int c : adj[v]) if (c != par[v]) {
            vector<ll> g = dp[c];
            g[0] = 1;
            vector<ll> nxt(cur.size() + g.size() - 1, 0);
            for (size_t a = 1; a < cur.size(); a++)
                for (size_t b = 0; b < g.size(); b++)
                    nxt[a + b] = (nxt[a + b] + cur[a] * g[b]) % MOD;
            cur.swap(nxt);
        }
        dp[v] = cur;
    }
    return dp[0];
}
 
int main() {
    int n = 8;
    vector<vector<int>> adj(n);
    auto edge = [&](int a, int b) { adj[a].push_back(b); adj[b].push_back(a); };
    edge(0, 1); edge(0, 2); edge(1, 3); edge(1, 4); edge(2, 5); edge(5, 6); edge(5, 7);
    vector<ll> w = count_root_sets(n, adj, 1000000007LL);
    for (int j = 1; j <= n; j++) printf("%lld ", w[j]);
    printf("\n");
}

Real output:

1 2 4 6 6 6 4 1 

The entry for size 1 is 1 (the root alone), and the entry for size 8 is 1 (the whole tree). The brute force over all subsets of this 8-node tree gives the same eight numbers.

Practice problems

Each problem has its own answer. The numbers come from brute force programs that I ran on the demo inputs. The nearest-marked answer was also checked on 2000 random small trees. The solution code was compiled with g++ -std=c++17 -O2.

  1. Maximum weight independent set. On the demo tree with skills 4, 7, 2, 9, 3, 6, 8 and 5, find the largest total skill of a set of people in which no two people are a manager and a report. Hint: each node has two states, chosen or not chosen. A chosen node forbids its children from being chosen. The state for each child is the best of the two states.

Show the answer

  1. The code below computes it, and a brute force over all 256 subsets also gives 29.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
 
ll mwis(int n, const vector<vector<int>> &adj, const vector<ll> &val) {
    vector<int> order, par(n, -1);
    vector<char> seen(n, 0);
    order.push_back(0); seen[0] = 1;
    for (size_t i = 0; i < order.size(); i++)
        for (int u : adj[order[i]]) if (!seen[u]) { seen[u] = 1; par[u] = order[i]; order.push_back(u); }
    vector<array<ll, 2>> dp(n, {0, 0});
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        dp[v][1] = val[v];
        for (int c : adj[v]) if (c != par[v]) {
            dp[v][0] += max(dp[c][0], dp[c][1]);
            dp[v][1] += dp[c][0];
        }
    }
    return max(dp[0][0], dp[0][1]);
}
  1. Nearest marked node. A tree has some marked nodes. For each node v, find the smallest weighted distance from v to any marked node. Use rerooting with the minimum instead of the maximum. Hint: down[v] is 0 if v is marked, and otherwise the best of down[c] + w. Start the up pass at the root with 0 if the root is marked and infinity otherwise.

Show the answer

The code below is the rerooting with the minimum. It was checked against a walk from every node on 2000 random trees with up to 9 nodes, with no mismatch.

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = LLONG_MAX / 4;
 
// For each node v, the smallest weighted distance from v to a marked node.
vector<ll> nearest_marked(int n, const vector<vector<pair<int, ll>>> &adj, const vector<int> &marked) {
    vector<int> order, par(n, -1);
    vector<ll> pw(n, 0);
    vector<char> seen(n, 0);
    order.push_back(0); seen[0] = 1;
    for (size_t i = 0; i < order.size(); i++)
        for (auto [u, w] : adj[order[i]]) if (!seen[u]) { seen[u] = 1; par[u] = order[i]; pw[u] = w; order.push_back(u); }
    vector<ll> down(n), up(n);
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        down[v] = marked[v] ? 0 : INF;
        for (auto [c, w] : adj[v]) if (c != par[v] && down[c] < INF) down[v] = min(down[v], down[c] + w);
    }
    up[0] = marked[0] ? 0 : INF;
    for (int i = 0; i < n; i++) {
        int v = order[i];
        for (auto [c, w] : adj[v]) if (c != par[v]) {
            ll best = marked[v] ? 0 : INF;
            if (up[v] < best) best = up[v];
            for (auto [c2, w2] : adj[v]) if (c2 != par[v] && c2 != c && down[c2] < INF) best = min(best, down[c2] + w2);
            up[c] = best < INF ? best + w : INF;
        }
    }
    vector<ll> res(n);
    for (int v = 0; v < n; v++) res[v] = min(down[v], up[v]);
    return res;
}

The code is a function, so a test or a main that builds a tree and marked nodes is needed to run it. The brute force for the same check is a walk from each node, as the lesson describes.

  1. Counting longest paths. In a DAG with weights, count the paths from the source that have the maximum total weight to each node. On the demo DAG, how many longest paths reach node 5? Hint: when a new path ties the best known distance, add the counts. When it beats it, replace the count.

Show the answer

  1. The longest path to node 5 is 0-1-4-5 with weight 13, and it is the only path that reaches that weight. The code below gives the same result as a list of all paths.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
 
pair<vector<ll>, vector<ll>> longest_and_ways(int n, const vector<vector<pair<int, ll>>> &adj, int src, ll MOD) {
    vector<int> indeg(n, 0);
    for (int v = 0; v < n; v++) for (auto [u, w] : adj[v]) indeg[u]++;
    vector<int> q, topo;
    for (int v = 0; v < n; v++) if (!indeg[v]) q.push_back(v);
    for (size_t i = 0; i < q.size(); i++) { int v = q[i]; topo.push_back(v); for (auto [u, w] : adj[v]) if (--indeg[u] == 0) q.push_back(u); }
    const ll NEG = LLONG_MIN / 4;
    vector<ll> best(n, NEG), ways(n, 0);
    best[src] = 0; ways[src] = 1;
    for (int v : topo) {
        if (best[v] == NEG) continue;
        for (auto [u, w] : adj[v]) {
            if (best[v] + w > best[u]) { best[u] = best[v] + w; ways[u] = ways[v]; }
            else if (best[v] + w == best[u]) ways[u] = (ways[u] + ways[v]) % MOD;
        }
    }
    return {best, ways};
}
  1. Merge counter. The lesson's knapsack counts its inner loop steps. For a star with a root and n - 1 leaves, predict the counter value for n = 10, then check your prediction by running the function. Hint: each leaf merge costs the current table size times 2, and the table grows by one size per leaf.

Show the answer

  1. The root starts alone, and the nine leaf merges happen with table sizes 1 to 9. Each merge costs the current size times 2, so the total is 2 times (1 + 2 + ... + 9), which is 2 times 45, or 90. The formula n times (n - 1) also gives 90. I ran the counter for n = 4, 5, 8 and 10, which gave 12, 20, 56 and 90.

Checklist

  • Build the order with a queue, not recursion, for any tree or graph that can be long.
  • Process a subtree after all of its children. Keep the parent's table separate from each child's table.
  • Use table sizes as loop bounds in a knapsack merge. This is what makes the total O(n^2).
  • For rerooting, write the down pass, then the up pass, and test the answer on every node against a walk from that node.
  • Keep the node itself in the candidate list of the up pass, with value 0.
  • For a DAG, check that the topological order has n nodes before you use any answer.
  • Use a modulus for counts and keep distances as plain sums.
  • Test every function on random small trees against a brute force that tries every subset or every path.

A prompt you can give to an AI tool for this kind of task:

Write a C++17 function that takes a weighted tree with n up to 200000 nodes and
returns, for every node, the largest weighted distance from that node to any
other node. Use rerooting with two passes and build the node order with a
queue, not recursion. Also write a brute force that runs a walk from every node.
Test both on 200 random trees with n from 1 to 9 and report any mismatch.

A good answer has a down pass, an up pass that keeps the top two child values, the node itself counted as 0, and no recursion. Check two things. The function must match the brute force on every node, and it must run on a path of 200000 nodes without a crash. If the answer uses recursion, the second check fails on long paths. If the up pass leaves out the node itself, the first check fails on the end nodes of a path.

Summary

For a tree, process children before parents and combine their tables. A knapsack merge gives the best value for each team size in O(n^2) total, because each pair of nodes is combined once. Rerooting gives an answer for every node in two passes. For a DAG, a topological order lets one pass compute longest paths and path counts. Build every order with a queue, check that a DAG really is acyclic, and test each function against a brute force on small random inputs.

Practice this topic: coding problems for DP optimisations →
← Digit DP