Dynamic programming on a tree

A tree hands you the independence an interval DP has to work for: the subtrees of different children never interact except through their parent.

The template

def solve(node, parent):
    results = [solve(child, node) for child in graph[node] if child != parent]
    return combine(node, results)

Passing the parent is what prevents walking back up an undirected edge. No visited set is needed, because a tree has no cycles.

Two values per node

The no two adjacent problem, on a tree:

def rob(node, parent):
    take = value[node]
    skip = 0
    for child in graph[node]:
        if child == parent: continue
        c_take, c_skip = rob(child, node)
        take += c_skip                  # taking me forbids my children
        skip += max(c_take, c_skip)     # skipping me frees them
    return take, skip

answer = max(rob(root, -1))

One value per node cannot work, because the best answer for a subtree depends on whether the parent was taken, which is information the subtree cannot see. That is the general signal for needing another dimension.

Return one thing, record another

Tree diameter needs two different quantities and conflating them is the usual mistake:

best = 0

def depth(node, parent):
    global best
    top_two = [0, 0]
    for child in graph[node]:
        if child == parent: continue
        d = depth(child, node)
        top_two = sorted(top_two + [d])[-2:]
    best = max(best, top_two[0] + top_two[1])   # a path through this node
    return 1 + top_two[1]                       # how far down I can reach

The returned value is how far down this node reaches, which the parent needs. The recorded value is the best path whose highest point is this node, which is useless to the parent because such a path cannot be extended upward.

Rerooting

When the answer is wanted for every node as the root, running the traversal n times is O(n^2). Two passes give O(n): one upward computing each subtree's answer, one downward passing each node the answer for everything outside its subtree.

The awkward part is that the parent must pass down its own value with the current child's contribution removed. For a sum that is subtraction; for a maximum it cannot be undone, so you keep the best two values and exclude one.

Watch the depth

A tree can be a chain, so recursion can be n frames deep. On 200,000 nodes that overflows the stack in most languages. Convert to an explicit stack or raise the limit, which is the same concern as any deep memoised recursion.