Trees, heaps and tries / Trees / 061 · Rerooting DP (sum over all roots)

Rerooting DP (sum over all roots)

Compute the answer for one root, then shift to each child in O(1) using subtree sizes.

The invariant

ans[child] = ans[parent] − size[child] + (n − size[child]).

01The intuition

Start with what you would do by hand

Six towns are joined by roads that form a tree: town 0 is joined to 1 and 2; town 2 is joined to 3, 4 and 5. For each town, we want the total distance to all the other towns. For town 0: distances 1, 1, 2, 2, 2, total 8. For town 2: 1, 1, 1, 1, 2, total 6. By hand you would stand in each town and walk out to every other town, adding up. Six towns, six walks of five steps each. For n towns that is n walks of n steps, n2; for 100,000 towns, ten billion steps.

But notice something between neighbours. Moving the starting town from 0 to its neighbour 2, every town in 2's group (2, 3, 4, 5) gets one step closer and every town outside it (0, 1) gets one step farther. So the new total is 8 − 4 + 2 = 6. One subtraction and one addition instead of a whole walk.

Build the structure, one piece at a time

Hang the tree from town 0 as the root. Now every town has a subtree: itself and everything below it.

Piece one: size[v], the number of towns in v's subtree. Computed bottom-up: a leaf has size 1; a parent adds its children's sizes plus 1. Here size[3] = size[4] = size[5] = size[1] = 1, size[2] = 4, size[0] = 6.

Piece two: down[v], the sum of distances from v to the towns in its own subtree. Also bottom-up: down[v] = sum over children c of (down[c] + size[c]), because every town under c is one step farther from v than from c. So down[2] = 0+1 + 0+1 + 0+1 = 3 and down[0] = (0 + 1) + (3 + 4) = 8. The root's down is its full answer, since its subtree is everything.

Piece three: reroot. Walk top-down. For a child c of p, the answer is the parent's answer shifted: the size[c] towns below c get closer by one, the other n − size[c] get farther by one.

Townsizedownanswer
0 (root)688
1108 − 1 + 5 = 12
2438 − 4 + 2 = 6
3106 − 1 + 5 = 10
4106 − 1 + 5 = 10
5106 − 1 + 5 = 10

The video uses a slightly different tree and so shows [8, 13, 6, 11, 11, 11]; the table here is for the six-town tree described above, which is easier to draw.

Why it is correct

Fix an edge between parent p and child c. Every town is either inside c's subtree or outside it. Standing at c instead of p, each inside town is exactly one step nearer and each outside town exactly one step farther, because the only road between the two groups is the edge p–c.

ans[c] = ans[p] − size[c] + (n − size[c])

The root's answer is correct from the bottom-up pass, and each child's answer follows from its parent's by the formula, so a top-down pass fills everything correctly.

Why it is fast

Pass one visits every town once to compute size and down. Pass two visits every town once to apply the formula. Two passes of n constant-time steps: O(n), against n2 for walking from every town. For 100,000 towns that is about 200,000 operations instead of ten billion.

How to recognise it in an interview

"For every node, compute something about the whole tree as seen from that node": sum of distances, number of nodes within reach, the best root to minimise height. Say: "I will compute the answer for one root bottom-up with subtree sizes, then reroot to each child in O(1)." Sum of Distances in Tree is exactly the formula above. Minimum Height Trees asks which roots minimise height; the same two-pass idea gives the height from every node, though peeling leaves layer by layer is the shorter solution there.

02The idea in one line

Compute the answer for one root, then shift to each child in O(1) using subtree sizes.

Reach for it when: Sum of distances in tree, best root choice, count paths through each node.

Cost: O(n) time.

03The walk-through, step by step

The video works one example: Sum of Distances in a 6-node tree. These are its captions in order, so you can follow without the video or check a step you missed.

  1. Pass 1, post-order: size[v] and the sum of distances inside v's subtree
  2. Node 1: size 1, subtree distance sum 0
  3. Node 3: size 1, subtree distance sum 0
  4. Node 4: size 1, subtree distance sum 0
  5. Node 5: size 1, subtree distance sum 0
  6. Node 2: size 4, subtree distance sum 3
  7. Node 0: size 6, subtree distance sum 8
  8. The root's answer is its own subtree sum: ans[0] = 8
  9. Move root to 1: 8 − 1 + (6 − 1) = 12
  10. Move root to 2: 8 − 4 + (6 − 4) = 6
  11. Move root to 3: 6 − 1 + (6 − 1) = 10
  12. Move root to 4: 6 − 1 + (6 − 1) = 10
  13. Move root to 5: 6 − 1 + (6 − 1) = 10
  14. Every node answered from its parent in O(1): [8,13,6,11,11,11]
  15. Two passes over the tree: O(n) instead of n separate traversals

04Asked at Google

3 Google-tagged LeetCode problems use this pattern. Frequency is LeetCode's Google-tag score (higher means reported more often); "last 6 months" marks problems reported in the most recent window of the export.

ProblemDifficultyFrequencyReported
Minimum Height TreesMedium33last 6 months
Sum of Distances in TreeHard31last 6 months
Time Needed to Inform All EmployeesMedium26

Other patterns in Trees: Recursive DFS (post-order fold), Best path through each node, Level-order BFS, Root-to-leaf path carry, Lowest common ancestor, Build from traversals, Iterative traversal with a stack, Column / coordinate traversal, Post-order deletion / pruning, Tree as a graph (parent pointers), Tree DP: take or skip this node, Serialize / deserialize by BFS.