DSU on Tree (Small to Large)

Answer subtree queries offline in $O(n\log n)$ without any range data structure. Despite the name it has nothing to do with DSU — it is the small-to-large merging technique applied to a tree.

The idea

Process each vertex’s children, keeping a global “counter” structure of what is in the current subtree. The trick:

• process all light children first, clearing the structure after each;
• process the heavy child last and keep its data;
• then re-add the light subtrees and the vertex itself, and answer the query.

Because each vertex is re-added only when it is inside a light subtree, and a vertex is inside at most $O(\log n)$ light subtrees on its root path, the total work is $O(n\log n)$.

Implementation

int cnt[MAXV];                 // the global counter structure
long long curAns;
vector<int> heavy, sz;
 
void add(int v, int p, int delta) {
    cnt[colour[v]] += delta;
    // maintain curAns incrementally here
    for (int c : adj[v])
        if (c != p && c != heavy[v] * KEEP) add(c, v, delta);
}
 
void dfs(int v, int p, bool keep) {
    // 1. light children first, discarding their contribution
    for (int c : adj[v])
        if (c != p && c != heavy[v]) dfs(c, v, false);
    // 2. heavy child last, keeping it
    if (heavy[v] != -1) dfs(heavy[v], v, true);
    // 3. add v and all light subtrees back
    add(v, p, +1, /*skip heavy*/ true);
    ans[v] = curAns;
    // 4. if this subtree's data is not needed by the parent, clear it
    if (!keep) add(v, p, -1, false);
}

The exact shape varies; what matters is: light children are recomputed, the heavy child is inherited.

Why $O(n\log n)$

A vertex $u$ is re-added once for each light edge on the path from $u$ to the root. Crossing a light edge at least halves the subtree size, so there are at most ${\log }_{2}n$ of them. Total additions: $O(n\log n)$.

What it answers

Anything expressible as an incrementally maintainable aggregate over the multiset of a subtree:

Query	Structure
Number of distinct colours in each subtree	a count array plus a distinct counter
Most frequent colour in each subtree	count array + count-of-counts
Sum of the $k$ most frequent	count array + a Fenwick tree over frequencies
Number of vertices at each depth	array indexed by depth
Number of pairs with a given property	maintain the pair count incrementally
Sum of subtree values matching a predicate	any incremental aggregate

The requirement is that add(x) and remove(x) each update the answer in $O(1)$ or $O(\log n)$ — the same requirement Mo’s algorithm imposes.

The alternatives

Approach	Time	Notes
DSU on tree	$O(n\log n)$	offline; simple; no heavy machinery
Small-to-large with set/map	$O(n{\log }^{2}n)$	merges arbitrary containers; even simpler to write
Euler tour + Mo’s algorithm	$O(n\sqrt{n})$	offline; works for path queries too
Euler tour + persistent segment tree	$O(n\log n)$	online; more memory
Merging segment trees	$O(n\log n)$	online; supports updates
HLD + segment tree	$O({\log }^{2}n)$ per query	for path, not subtree, queries

Small-to-large with containers

The lazier variant: give each vertex a set or map, and merge children into the largest one.

void dfs(int v, int p) {
    for (int c : adj[v]) {
        if (c == p) continue;
        dfs(c, v);
        if (s[v].size() < s[c].size()) swap(s[v], s[c]);
        for (auto x : s[c]) s[v].insert(x);
        s[c].clear();
    }
    s[v].insert(colour[v]);
    ans[v] = s[v].size();
}

Ten lines, $O(n{\log }^{2}n)$, and it handles “number of distinct values per subtree” immediately. Reach for this first; upgrade to the heavy-child version only if the log factor matters.

See also: HLD · Mo’s Algorithm · Euler Tour