Unit 13: Trees

Pick a node and let the rest hang

A tree has no cycles, so nothing you do can go round in a circle. Choose any node as the root and every other node gets a parent — the neighbour one step closer to the root — and a depth.

One search fills in both.

Python

n = 7
edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (5, 6)]
adj = [[] for i in range(n)]
for a, b in edges:
    adj[a].append(b)
    adj[b].append(a)

parent = [-1] * n
depth = [0] * n

def walk(node, from_node):
    parent[node] = from_node
    for nxt in adj[node]:
        if nxt != from_node:
            depth[nxt] = depth[node] + 1
            walk(nxt, node)

walk(0, -1)
print(parent)
print(depth)

It prints

[-1, 0, 0, 1, 1, 2, 5]
[0, 1, 1, 2, 2, 2, 3]

No visited list needed

In a general graph you carry a visited array. In a tree, "do not go back the way you came" is enough — there is no other way to meet an old node.

That is the whole reason tree problems are easier than graph problems, and it is worth saying out loud once.

The tree in the sample

Node 0 at the top, with 1 and 2 below it. Under 1 are 3 and 4; under 2 is 5, and under 5 is 6.

Depths: 0 is 0, then 1 and 2 are 1, then 3, 4 and 5 are 2, and 6 is deepest at 3. Keep it in your head — the rest of the unit uses this tree.

How big is each branch?

The subtree of a node is that node plus everything hanging below it. Its size is 1, plus the sizes of its children — which you only know once you have been down there.

So the work happens on the way back up: recurse first, add afterwards.

Python

n = 7
adj = [[1, 2], [0, 3, 4], [0, 5], [1], [1], [2, 6], [5]]
size = [0] * n

def walk(node, from_node):
    size[node] = 1
    for nxt in adj[node]:
        if nxt != from_node:
            walk(nxt, node)
            size[node] += size[nxt]

walk(0, -1)
print(size)

It prints

[7, 3, 3, 1, 1, 2, 1]

Anything that adds up works the same way

Swap the count for a sum of values and you have the total in each subtree. Swap it for a max and you have the deepest node below each one.

This is dynamic programming again — the state is a node, the moves are its children, and the order is "children before parents", which the recursion gives you free.

Python

n = 7
adj = [[1, 2], [0, 3, 4], [0, 5], [1], [1], [2, 6], [5]]
value = [5, 2, 1, 4, 3, 6, 7]
total = [0] * n

def walk(node, from_node):
    total[node] = value[node]
    for nxt in adj[node]:
        if nxt != from_node:
            walk(nxt, node)
            total[node] += total[nxt]

walk(0, -1)
print(total)

It prints

[28, 9, 14, 4, 3, 13, 7]

Sizes answer questions that look much harder

"How many pairs of nodes have their route passing through this edge?" — the size below it, times everyone else. For the edge above node 1 that is 3 times 4, which is 12.

No pair of nodes is ever listed. The counting is done by the sizes.

Python

n = 7
size = [7, 3, 3, 1, 1, 2, 1]
parent = [-1, 0, 0, 1, 1, 2, 5]

for node in range(1, n):
    print(node, parent[node], size[node] * (n - size[node]))

It prints

1 0 12
2 0 12
3 1 6
4 1 6
5 2 10
6 5 6

How far apart are the two furthest nodes?

That distance is the tree's diameter. Checking every pair costs n * n searches, which is far too slow — and there is a two-search trick instead.

Search from anywhere and find the furthest node. Search again from there. That second distance is the diameter.

Two searches, and that is all

From node 0 the furthest is 6, three steps away. From 6 the furthest is 3 or 4, five steps away — so the diameter is 5.

The route is 3 to 1 to 0 to 2 to 5 to 6: five edges, and you can check it by eye on the picture from lesson 1.

Python

from collections import deque

n = 7
adj = [[1, 2], [0, 3, 4], [0, 5], [1], [1], [2, 6], [5]]

def far_from(start):
    dist = [-1] * n
    dist[start] = 0
    queue = deque([start])
    best = start
    while queue:
        node = queue.popleft()
        if dist[node] > dist[best]:
            best = node
        for nxt in adj[node]:
            if dist[nxt] == -1:
                dist[nxt] = dist[node] + 1
                queue.append(nxt)
    return best, dist[best]

end, d = far_from(0)
print(end, d)
other, diameter = far_from(end)
print(other, diameter)

It prints

6 3
3 5

Why the furthest node is always an end of a longest path

Suppose the furthest node from your start is not on any longest path. Follow where the start's route joins that longest path and compare the two branches — the one to your furthest node is at least as long, so swapping it in gives a path at least as long.

So it is an end of one. That is why the second search finds the whole thing.

Where do two nodes meet?

Walk up from 3 and up from 6 until the two walks land on the same node. That node is the lowest common ancestor, and here it is 0.

It is also how you get the distance between any two nodes: depth[a] + depth[b] - 2 * depth[meeting point].

The slow way first

Level them up — walk the deeper one until the depths match — then step both up together until they touch.

Correct, and O(n) per question. A problem with 200000 questions needs it faster.

Python

parent = [-1, 0, 0, 1, 1, 2, 5]
depth = [0, 1, 1, 2, 2, 2, 3]

def meet(a, b):
    while depth[a] > depth[b]:
        a = parent[a]
    while depth[b] > depth[a]:
        b = parent[b]
    while a != b:
        a = parent[a]
        b = parent[b]
    return a

print(meet(3, 4), meet(3, 6), meet(6, 5), meet(4, 4))

It prints

1 0 5 4

Unit 10 already built the fast way

A successor graph where each node points at its parent — that is exactly what a rooted tree is. So build the same doubling table: up[k][node] is the ancestor 2 ** k steps above.

The root points at itself so a jump can never fall off the top.

Python

n = 7
parent = [0, 0, 0, 1, 1, 2, 5]
levels = 3

up = [parent[:]]
for k in range(1, levels):
    prev = up[-1]
    up.append([prev[prev[i]] for i in range(n)])

for k in range(levels):
    print(2 ** k, up[k])

It prints

1 [0, 0, 0, 1, 1, 2, 5]
2 [0, 0, 0, 0, 0, 0, 2]
4 [0, 0, 0, 0, 0, 0, 0]

Level up by bits, then rise together

To lift a node by k levels, jump by the bits of k. Then, if the two are still different, jump both up by the biggest amount that keeps them apart — because landing on the same node might overshoot the lowest one.

When no jump can be taken any more, one more step gets there.

Python

n = 7
parent = [0, 0, 0, 1, 1, 2, 5]
depth = [0, 1, 1, 2, 2, 2, 3]
levels = 3

up = [parent[:]]
for k in range(1, levels):
    prev = up[-1]
    up.append([prev[prev[i]] for i in range(n)])

def lift(node, steps):
    k = 0
    while steps > 0:
        if steps & 1:
            node = up[k][node]
        steps >>= 1
        k += 1
    return node

def lca(a, b):
    if depth[a] < depth[b]:
        a, b = b, a
    a = lift(a, depth[a] - depth[b])
    if a == b:
        return a
    for k in range(levels - 1, -1, -1):
        if up[k][a] != up[k][b]:
            a = up[k][a]
            b = up[k][b]
    return parent[a]

print(lca(3, 4), lca(3, 6), lca(6, 5), lca(4, 2))

It prints

1 0 5 0