Small-to-Large Merging, in depth: the doubling argument, subtree queries in O(n log n), carrying aggregates, the sack variant, and where it breaks

The problem: merging collections repeatedly

Take a rooted tree with n nodes, each with a colour, and ask: for every node, how many distinct colours appear in its subtree? A direct solution builds each node's set from scratch, which costs the subtree size per node and so O(n squared) on a path. A smarter-looking solution builds each node's set by merging its children's sets into a new one, but copying is still expensive: on a path of n nodes the bottom set is copied n - 1 times, the next n - 2 times and so on, again O(n squared).

The waste is moving the big set. When a node's child already holds a set of 10,000 colours and the node itself adds one, the right move is to add one element to the child's set and adopt it, not to copy 10,000 elements. Small-to-large generalises that observation: at every merge, compare sizes, move the smaller collection's elements into the larger, and give the larger one to the parent by reference. On a path this makes the total cost n - 1 moves instead of roughly n squared over 2.

Why each element moves at most log n times

Start with the case where all elements are different, and follow one element. Every time it is moved, it was in the smaller of two collections, so the collection it lands in has at least twice as many elements as the one it left. A collection can never hold more than n elements, so an element can be moved at most log2 n times. Summing over n elements gives at most n log2 n moves in total, regardless of the shape of the tree or the order of merges.

Duplicates need one more step, because a set that drops repeated colours may not double. Give every collection a weight equal to the number of tree nodes merged into it, so weights add exactly. A merge costs the smaller size, and each size is at most its weight, so the cost is at most the smaller weight. Charge that cost to the nodes on the lighter side: each such node's group at least doubles in weight, which can happen at most log2 n times per node. The n log2 n bound therefore holds for deduplicated sets too; deduplication only makes merges cheaper.

Total running time is the number of moves times the cost of one insertion. With hash sets and dictionaries, insertion is expected O(1), so the whole algorithm is expected O(n log n). With balanced-tree containers such as C++ std::set or Java TreeMap, insertion is O(log n), so the total is O(n log squared n).

The bound is tight up to a constant. On a complete binary tree every merge is between equal halves, so each element moves at every level above it. We measured the code below on a complete binary tree of 16,384 nodes with a distinct value per node: 106,497 moves, close to n log2 n / 2 = 114,688. On a path of the same size it made 16,383 moves, one per node.

Implementation: distinct colours in every subtree

The implementation needs a children-before-parents order. Recursion is natural, but Python's default recursion limit is 1,000, so a tree shaped like a path breaks it; an explicit stack is safer. Reversing a depth-first preorder lists every node after all of its descendants, which is all the merge loop needs.

def preorder(n, edges, root=0):
    adj = [[] for _ in range(n)]
    for a, b in edges:
        adj[a].append(b)
        adj[b].append(a)
    parent, order, stack = [-1] * n, [], [root]
    seen = [False] * n
    seen[root] = True
    while stack:
        u = stack.pop()
        order.append(u)
        for v in adj[u]:
            if not seen[v]:
                seen[v] = True
                parent[v] = u
                stack.append(v)
    return parent, order

def distinct_colours(n, edges, colour, root=0):
    parent, order = preorder(n, edges, root)
    sets = [{colour[u]} for u in range(n)]
    ans = [0] * n
    for u in reversed(order):          # every child before its parent
        ans[u] = len(sets[u])          # u's set is complete: read it now
        p = parent[u]
        if p >= 0:
            big, small = sets[p], sets[u]
            if len(big) < len(small):
                big, small = small, big    # swap references, not contents
            big.update(small)
            sets[p] = big
            sets[u] = None             # u's set may now belong to p
    return ans

Two lines carry the correctness. The answer for u is read before its set is handed to the parent, because afterwards that same object keeps growing with the parent's other children. And the swap exchanges references: in Python that is reassigning names, in C++ it is std::swap on the containers, which is constant time. Copying, or passing a container by value in C++, quietly reintroduces the quadratic cost.

Carrying aggregates through the merge

Sets answer membership questions; most real problems need an aggregate. A classic example: for each subtree, sum the colours that occur most often (several colours can tie). The trick is to keep, next to each frequency table, the aggregate itself, and update it only on insertions into the big table. Counts only ever increase during a merge, so the maximum frequency can only rise, and every change to it happens at an insertion we already pay for.

def dominant_colour_sum(n, edges, colour, root=0):
    parent, order = preorder(n, edges, root)
    # bag = [counts, max_frequency, sum_of_colours_with_that_frequency]
    bags = [[{colour[u]: 1}, 1, colour[u]] for u in range(n)]
    ans = [0] * n
    for u in reversed(order):
        ans[u] = bags[u][2]
        p = parent[u]
        if p < 0:
            continue
        big, small = bags[p], bags[u]
        if len(big[0]) < len(small[0]):
            big, small = small, big
        counts = big[0]
        for col, k in small[0].items():
            f = counts.get(col, 0) + k
            counts[col] = f
            if f > big[1]:
                big[1], big[2] = f, col
            elif f == big[1]:
                big[2] += col
        bags[p], bags[u] = big, None
    return ans

Size is measured by the number of distinct keys, which is the number of moves a merge costs, and the moved unit is a (colour, count) pair, so the doubling argument applies unchanged. Both functions in this article were checked against a brute-force subtree scan on hundreds of random trees. The general rule: an aggregate is cheap to carry if it can be updated from insertions alone. Maximum frequency, sum, count of distinct values and the answer to a threshold query can; a median or an arbitrary order statistic needs an ordered container and brings back the extra log factor.

Worked example: a seven-node tree

Trace the seven-node tree in the diagram: node 0 (red) has children 1 (blue) and 2 (red); node 1 has children 3 (green) and 4 (blue); node 2 has child 5 (red); node 4 has child 6 (green). The iterative preorder visits 0, 2, 5, 1, 4, 6, 3, so the loop processes 3, 6, 4, 1, 5, 2, 0.

Seven moves in total, against eight if every child's set were always merged into its parent's regardless of size. On a tree this small the saving is one move; on a path it is the difference between n - 1 moves and roughly n squared over 2. The interesting row is node 1: the root held only {red}, so rather than copying two colours up, the algorithm moved red down into node 1's set and gave that set to the root. Ties can go either way; the bound only needs the kept collection to be at least as large as the moved one.

Variants: merging groups, the sack, and relatives

Merging groups. Outside trees the same rule maintains explicit membership lists as groups merge: keep a list per group and a group id per element, and when two groups merge, relabel and append the smaller list. Every element is relabelled at most log2 n times, which gives O(n log n) total and lets you list a group's members at any time, something plain union-find cannot do.

The sack (DSU on tree). A variant avoids hash tables entirely. Compute subtree sizes and each node's heavy child (the child with the largest subtree), and lay out an Euler tour so each subtree is a contiguous range. Process light children first and clear their contributions from one global counter array; then process the heavy child and keep its contributions; then add the node and its light subtrees by scanning their ranges. A node is re-added once for each light edge between it and the root, and any path to the root crosses at most log2 n light edges, so this is also O(n log n), with flat arrays that are fast in practice. Its shape:

def dfs(u, keep):
    for v in light_children(u):
        dfs(v, keep=False)             # answer v, then wipe its counts
    if heavy[u] is not None:
        dfs(heavy[u], keep=True)       # its counts stay in the global array
    for v in light_children(u):
        add_range(tin[v], tin[v] + size[v])   # Euler-tour range of v
    add(u)
    answer[u] = query()                # global array now holds exactly u's subtree
    if not keep:
        remove_range(tin[u], tin[u] + size[u])

Other relatives. When queries are about subtrees and can be answered offline, flattening the tree with an Euler tour turns each subtree into an array range, and range techniques apply directly: a Fenwick or segment tree for sums and counts, or Mo's algorithm for distinct-value style questions. Segment tree merging, where each node owns a sparse segment tree over values and children's trees are merged recursively, gives O(n log n) total with order statistics included. Path queries rather than subtree queries usually need lowest common ancestor machinery instead.

Trade-offs and failure modes

Failure modes are mostly about accidentally breaking the rule. Copying instead of adopting, comparing the wrong sizes (for example list length when duplicates make the set much smaller), or always merging child into parent regardless of size all restore quadratic time on paths and caterpillar trees, and none of them fail on random tests, which tend to be shallow. Add a worst-case test: a path of 200,000 nodes and a complete binary tree, timed. Reading an answer after its collection has been adopted returns the parent's data. Problems keyed by depth (how many nodes at depth d in this subtree) cannot use plain dictionaries keyed by absolute depth without care; store relative depth with an offset, or use the long-path variant that inherits the deepest child's array. Finally, the technique is batch-oriented: a single merge can take time proportional to n, so it is a poor fit for latency-sensitive online systems unless merges are bounded.

What to do next

1. Implement distinct_colours with an explicit stack and check it against a brute-force subtree scan on random trees.
2. Time it on a 200,000-node path and on a complete binary tree, then remove the size comparison and watch the path case collapse.
3. Add an aggregate, such as the dominant-colour sum, that updates only on insertion, and test ties.
4. Rewrite the solution as a sack with a global counter array and compare speed and memory.
5. Use the same rule outside trees: maintain explicit member lists for merging groups and list a group on demand.
6. Learn the Euler-tour flattening so you can switch to range data structures when queries are offline.