DP on trees

A tree has no cycles, so it splits cleanly at every vertex: what happens below a vertex does not depend on anything outside, except through that one vertex. That makes trees the easiest place to do dynamic programming. Each vertex gets an answer for the part of the tree below it, built from its children’s answers.

A state for each subtree

Pick any vertex as the top of the tree and call it the root (TREE-1). Every other vertex then has one parent, the next vertex on the way up, and its children are the vertices just below it. The subtree of v is v together with everything below it.

Two subtrees of different children of v share no vertex, and no edge joins them: every edge from inside a child’s subtree to the rest of the tree goes through that child and v. So a question about v’s subtree usually splits into the same question for each child’s subtree, plus a decision about v itself. That is the DP state: one answer (or a few) per subtree.

The order of work follows from it. A vertex needs its children’s answers first, so the vertices are handled children first, root last. This is a post-order (TREE-2). There are n states and each child is read once by its parent, so tree DPs usually cost O(n).

Take or skip

A line of stations forms a tree of roads. Each station has a value, and you want to staff some of them, but never two stations joined by a road. What is the largest total value? On a general graph this is the maximum weight independent set, an NP-hard problem (HARD-1). On a tree it takes one pass.

One answer per subtree is not enough, because the parent needs to know whether the child itself was staffed. So keep two:

• take(v): the best total in v’s subtree when v is staffed. Then no child may be staffed, so add the value of v and each child’s skip.
• skip(v): the best total when v is not staffed. Then each child is free, so add the better of each child’s two answers.

Why is that right? A choice in v’s subtree is a choice of v plus a separate choice in each child’s subtree. The only rule that links them is the road between v and a child. Once v’s own choice is fixed, the children no longer affect each other, so each can be made as good as possible on its own. The answer is the better of the root’s two values. Step through the figure; then switch it to rerooting, which is the last section.

def best_staffing(children, value, root=0):
    def solve(v):                 # (take, skip) for v's subtree
        take, skip = value[v], 0
        for c in children[v]:
            t, s = solve(c)
            take += s             # v staffed: c is not
            skip += max(t, s)     # v not staffed: c is free
        return take, skip
    return max(solve(root))

To list the staffed stations, walk back down from the root, as the figure’s last step does: if the root takes, its children skip; a child that may choose takes its better answer, and so on down.

Without recursion

The recursive version is short, but its depth is the height of the tree. A path of 100,000 stations is a tree too, and Python stops at about 1,000 nested calls (RECR-2). In the browser the limit is lower still. So for big trees, make the order by hand:

1. Go through the tree from the root in BFS order (GRPH-2). Each vertex is listed after its parent, and you note each vertex’s parent.
2. Go through that list backwards. Now each vertex comes after all of its children: a valid children-first order.

def best_staffing(n, roads, value):
    adj = [[] for _ in range(n)]
    for u, v in roads:
        adj[u].append(v)
        adj[v].append(u)
    parent = [-1] * n
    order = [0]
    for v in order:                   # BFS: each parent before its children
        for c in adj[v]:
            if c != parent[v]:
                parent[c] = v
                order.append(c)
    take = value[:]
    skip = [0] * n
    for v in reversed(order):         # each child before its parent
        p = parent[v]
        if p >= 0:
            take[p] += skip[v]
            skip[p] += max(take[v], skip[v])
    return max(take[0], skip[0])

Two details. The for v in order loop sees the items appended during the loop, so it walks the whole tree. And the test c != parent[v] stops the search going back up the road it came down; in a tree that is the only way back to a vertex already listed. Instead of the parent asking each child, each child here pushes its answers up to its parent. The sums come out the same.

Other answers per subtree

• Sizes. size(v) = 1 + size of each child: the number of vertices in each subtree, in one backward pass.
• The longest route (the diameter). Keep deep(v), the longest way down from v. The longest route that has v as its highest vertex goes down through v’s two deepest children, so it is the sum of the two largest deep(c) + 1 values. The diameter is the largest of these over all v.
• The fewest guards on roads. To cover every road with a guard at one of its ends (a vertex cover, FLOW-5), keep with(v) and without(v). If v has no guard, every child must have one.
• A budget spread over a tree. When the state is (v, how many vertices chosen inside v’s subtree), a parent joins its children’s tables one at a time. Joining tables of sizes a and b costs about a × b. Count it as pairs of vertices: each pair (x, y) is paid for once, at the vertex where their subtrees first meet. So the whole tree costs O(n²) (n · n), not the O(n³) (n · n · n) that a quick count suggests.

Every station as the top

Some questions ask for an answer at every vertex as if it were the root. For example: for each station, the total distance to all the other stations. Running the one-root DP from every vertex costs n × O(n) = O(n²). Rerooting does it in O(n) with two passes:

1. Up (children first), with root r: size(v), and down(v), the total distance from v to the vertices of its subtree. Each vertex in a child c’s subtree is one road further from v than from c, so down(v) adds down(c) + size(c) over the children.
2. Down (parents first): total(r) = down(r). Then move the answer from each parent p to each child c.

Why. Take the distances from p and change the start to c, one road further down. The size(c) vertices in c’s subtree are reached from c without the road p–c, so each is exactly 1 closer. Every other vertex, n − size(c) of them, is reached from c by first going up to p, so each is exactly 1 further. No vertex is in both groups. With roads of different lengths, the change is the length of the road p–c instead of 1, which is the practice problem.

The pattern is general. Pass 1 computes each answer looking down. Pass 2 gives each child the answer for “everything that is not in my subtree”, which the parent can build from its own answer and its other children. When the combining step cannot be undone (a maximum rather than a sum), keep the best and second-best child values, so that a child that gave the best can be given the second-best.

Now try ADDP-P1, Total distance from every station.