DSU on Tree (Small-to-Large), in depth: the keep-heavy sack, an iterative layout, offline subtree queries and the bugs that make it quadratic

Keep the heavy child, re-add the rest

Root the tree and compute every subtree size. Each vertex's heavy child is the child with the largest subtree; the edge to it is heavy and every other child edge is light. Any path from a vertex to the root crosses at most log2 n light edges, because the subtree size at least doubles at each one. This is the same decomposition that heavy-light decomposition uses for path queries.

The sack processes vertices bottom-up while keeping a single global state. At vertex u, it wants the state to describe exactly u's subtree. The heavy child's subtree is the largest part, so the algorithm arranges for it to still be in the state when u is reached. It then adds u itself and every vertex of every light child's subtree one at a time, records the answer, and, if u is itself a light child, wipes everything it added so that u's siblings start from an empty state.

A vertex is added once as itself and once more for every light edge between it and the root, when an ancestor re-adds the light subtree containing it. That is at most 1 + log2 n additions per vertex and the same number of wipes, so the total is O(n log n) operations, each O(1) for array counters.

The state never has to support removing one vertex while others remain. It only grows, and is then reset to empty in one sweep. That matters, because many aggregates cannot be maintained under deletion: the maximum frequency can drop when a vertex leaves, and recovering the new maximum would need another structure. Under add-and-wipe, a running maximum is enough.

A layout without recursion

The textbook version is a recursive DFS with a keep flag. That recursion is dangerous: a path-shaped tree of 200,000 vertices overflows Python's stack and can overflow a C++ thread stack too. The iterative layout below exploits one observation. If each vertex's children are visited with the heavy child last, then in the resulting post-order the heavy child is the last vertex finished before its parent, so its kept state is exactly what the parent finds. Light siblings finish earlier and wipe themselves before the next sibling starts.

The same traversal produces an Euler tour, in which every subtree is a contiguous range of positions from tin[u] to tin[u] + size[u] - 1, so a light subtree is re-added or wiped by scanning a slice of an array rather than walking pointers. Euler tours themselves are covered in the Euler tour technique.

Implementation: dominant colours in every subtree

The example query: for every vertex, the sum of all colours that occur the maximum number of times in its subtree. It needs a count per colour, the current maximum frequency and the running sum, which is exactly the kind of aggregate that deletion would break.

def dsu_on_tree(n, adj, colour, root=0):
    parent = [-1] * n
    order = [root]
    parent[root] = root
    for u in order:                          # BFS: parents before children
        for v in adj[u]:
            if v != parent[u]:
                parent[v] = u
                order.append(v)
    parent[root] = -1
    size, heavy = [1] * n, [-1] * n
    for u in reversed(order):                # children before parents
        if parent[u] >= 0:
            size[parent[u]] += size[u]
    for u in order:
        best = 0
        for v in adj[u]:
            if v != parent[u] and size[v] > best:
                best, heavy[u] = size[v], v

    tin, euler, post = [0] * n, [], []
    stack = [(root, False)]
    while stack:                             # Euler tour, heavy child visited last
        u, done = stack.pop()
        if done:
            post.append(u)
            continue
        tin[u] = len(euler)
        euler.append(u)
        stack.append((u, True))
        if heavy[u] >= 0:
            stack.append((heavy[u], False))  # pushed first, popped last
        for v in adj[u]:
            if v != parent[u] and v != heavy[u]:
                stack.append((v, False))

    cnt = [0] * (max(colour) + 1)
    maxf = best_sum = 0
    ans = [0] * n

    def add(x):
        nonlocal maxf, best_sum
        c = colour[x]
        cnt[c] += 1
        if cnt[c] > maxf:
            maxf, best_sum = cnt[c], c
        elif cnt[c] == maxf:
            best_sum += c

    for u in post:                           # state == heavy[u]'s subtree here
        add(u)
        for v in adj[u]:
            if v != parent[u] and v != heavy[u]:
                for i in range(tin[v], tin[v] + size[v]):
                    add(euler[i])
        ans[u] = best_sum
        p = parent[u]
        if p < 0 or heavy[p] != u:           # light child: wipe by scanning its range
            for i in range(tin[u], tin[u] + size[u]):
                cnt[colour[euler[i]]] = 0
            maxf = best_sum = 0
    return ans

This exact function was checked against a brute-force solver on 300 random trees of up to 40 vertices with up to five colours, and every answer matched.

Worked example: nine vertices by hand

Trace the post-order, which is v3, v6, v2, v4, v8, v7, v5, v1, v0. v3 is a light leaf: add colour 2, answer 2, wipe. v6 is the heavy child of v2, so after answering 3 it stays. At v2 the state already holds v6; adding v2 makes colour 3 occur twice, answer 3. v2 is light under v0, so it wipes its two positions. v4 and v8 are light leaves and wipe themselves. v7 is heavy under v5 and stays. At v5 the algorithm adds v5 and re-adds v8: colours 1, 2 and 3 once each, answer 1 + 2 + 3 = 6. v5 is heavy under v1 and stays.

At v1 the algorithm adds v1 (colour 2) and re-adds v4 (colour 1): colours 1 and 2 now occur twice and colour 3 once, so the answer is 3. Finally at v0, with v1's five vertices kept, it adds v0 and re-adds the light subtrees of v2 and v3. Every colour occurs three times and the answer is 6. The output is 6, 3, 3, 2, 1, 6, 3, 1, 3 for v0 to v8, matching the brute force. v8 was added twice: once as itself and once when v5 re-added its light subtree. It lies under one light edge, and the kept heavy subtrees at v1 and v0 cost nothing extra.

Other query shapes

Any query that can be answered from an add-only aggregate fits the same skeleton. Only add, the wipe, and the line that reads the answer change.

Distinct values. Keep a count per value and a counter of values whose count is non-zero.

Queries with parameters. If the questions arrive as pairs (v, k), attach each to its vertex beforehand and answer all of v's questions at the moment the state equals v's subtree. Everything is offline; if answers must be given before later questions are known, this technique does not apply.

Thresholds. To answer questions (v, k) of the form 'how many colours occur at least k times in v's subtree', keep a second array at_least, where at_least[f] counts the colours whose count has reached f. Each addition raises one colour's count by one and increments exactly one entry, so the answer is a single lookup. Here the wipe undoes each addition in turn, which is still proportional to the range:

def add(x):
    c = colour[x]
    cnt[c] += 1
    at_least[cnt[c]] += 1          # this colour now counts towards thresholds 1..cnt[c]

def wipe(lo, hi):
    for i in range(lo, hi):
        c = colour[euler[i]]
        at_least[cnt[c]] -= 1
        cnt[c] -= 1

# at vertex u: for (k, qid) in questions[u]: out[qid] = at_least[k] if k <= n else 0

Per-depth questions. For questions of the form 'how many distinct colours occur among vertices at depth d in v's subtree', index the state by depth:

depth_cnt = defaultdict(Counter)   # depth -> colour -> count
distinct = defaultdict(int)        # depth -> number of colours with count > 0

def add(x):
    d, c = depth[x], colour[x]
    depth_cnt[d][c] += 1
    if depth_cnt[d][c] == 1:
        distinct[d] += 1

def wipe(lo, hi):                  # Euler-tour range of a light subtree
    for i in range(lo, hi):
        d = depth[euler[i]]
        depth_cnt[d].clear()
        distinct[d] = 0

# at vertex u: for (d, qid) in questions[u]: out[qid] = distinct[d]

The wipe clears only the depths that appear in the range, never the whole table, which keeps it proportional to the subtree size.

What it costs in practice

Counting additions makes the bound concrete. On a 100,000-vertex path, every child is heavy, and the sack performs exactly 100,000 additions. A version that keeps no child and rebuilds every subtree performs 5,000,050,000. On a random recursive tree of the same size the sack performs 479,543 additions, 4.8 per vertex, against 1,302,070 for the rebuild-everything version. The log2 n bound is a ceiling, and on typical trees the count is far below it.

Timed in CPython 3.13, best of three runs, the function above handles a 200,000-vertex path in 0.29 seconds, a random recursive tree of the same size in 2.5 seconds and a complete binary tree in 2.4 seconds, with no recursion limit to raise. The path is fast because nothing is ever re-added. Bushy, balanced trees are the slow case, because most vertices sit under several light edges.

Failure modes

Wiping the whole array. Resetting cnt = [0] * K after each light child costs O(K) per light vertex, O(nK) in total. Wipe only the range you added.

Keeping the wrong child. Choosing the first child instead of the largest, or computing sizes before the parent array is final, breaks the doubling argument. On a caterpillar whose spine vertices list a leaf before the next spine vertex, keeping the first child re-adds the whole spine at every level: quadratic.

Forgetting the scalars. The wipe must also reset the derived values: maximum frequency, running sums and distinct counters. Missing one produces answers that are right for the first light subtree and wrong after it.

Answering at the wrong moment. Read the answer after the light subtrees are re-added and before the wipe; reading it earlier misses them.

Recursion depth. Recursive versions fail on long paths with stack overflows, which in contests show up as runtime errors on a single test.

Trade-offs against other techniques

Mo's algorithm, the most common alternative for distinct-value queries, is covered in Mo's algorithm for offline queries.

What to do next

1. Copy the function above and run it, with a brute-force checker, on random trees until you trust it.
2. Rewrite add and the answer line for distinct-colour counting, then for the per-depth question.
3. Feed it a 200,000-vertex path and confirm it neither recurses nor slows down.
4. Count additions on your real input; if the count per vertex approaches log2 n, check how heavy children are chosen.
5. If your queries need updates between them, switch to an Euler tour with a Fenwick or segment tree instead.