vii. Small to Large (DSU on Tree)

Answer a query about every subtree by keeping one global structure: process the light children and throw their data away, process the heavy child and keep its data, then re-insert the light subtrees and the vertex. Every vertex is re-inserted once per light edge above it, so there are $O(n\log n)$ insertions. The container version (merge the smaller set into the larger) is shorter and costs one more log. Despite the common name “DSU on tree,” this has nothing to do with DSU — it’s small-to-large merging applied to a tree; the requirement throughout is that add/remove each update the running answer in $O(1)$ or $O(\log n)$, the same requirement Mo’s algorithm imposes.

Approach	Time	Notes
DSU on tree (sack)	$O(n\log n\cdot c)$ for $O(c)$ add / remove	offline subtree queries, add/remove with $O(1)$ answer update
merge containers small into large	$O(n{\log }^{2}n)$ with set / map	simplest; also merges heaps, hash maps, segment trees — reach for this first, upgrade to the heavy-child version only if the extra log matters
Euler tour + persistent / merging segment trees	$O(n\log n)$	online
Euler tour + Mo	$O(n\sqrt{n})$	when add/remove is all you have and queries are general intervals
HLD + segment tree	$O({\log }^{2}n)$ per query	wrong tool here — that’s for path, not subtree, queries

Other queries this same machinery answers — anything expressible as an incrementally maintainable aggregate over a subtree’s multiset:

Query	Structure
number of distinct colours in each subtree	a count array plus a distinct counter
most frequent colour / its sum	count array + count-of-counts (as in dominatingColourSums below)
sum of the $k$ most frequent values	count array + Fenwick tree over frequencies
number of vertices at each depth within a subtree	array indexed by depth
counts of pairs with some property	maintain the pair count incrementally alongside the per-value counts
sum of subtree values matching a predicate	any incremental aggregate

DSU on tree

struct DsuOnTree {
    int n; vector<vector<int>> adj; vector<int> par, sub, heavy, tin, tout, ver;
    DsuOnTree(const vector<vector<int>>& g, int root) : n(sz(g)), adj(g), par(n, -1), sub(n, 1), heavy(n, -1), tin(n), tout(n), ver(n) {
        vector<int> order{root};
        for (int i = 0; i < n; i++) for (int v : adj[order[i]]) if (v != par[order[i]]) par[v] = order[i], order.push_back(v);
        for (int i = n - 1; i > 0; i--) sub[par[order[i]]] += sub[order[i]];
        for (int u = 0; u < n; u++) for (int v : adj[u]) if (v != par[u] && (heavy[u] < 0 || sub[v] > sub[heavy[u]])) heavy[u] = v;
        vector<int> st{root}; int t = 0;                     // preorder: subtree(v) = ver[tin[v] .. tout[v])
        while (sz(st)) { int u = st.back(); st.pop_back(); ver[tin[u] = t++] = u; tout[u] = tin[u] + sub[u]; for (int v : adj[u]) if (v != par[u]) st.push_back(v); }
    }
    template <class Add, class Rem, class Ans> void run(int root, Add add, Rem rem, Ans answer) {
        auto dfs = [&](auto& self, int v, bool keep) -> void {
            for (int c : adj[v]) if (c != par[v] && c != heavy[v]) self(self, c, false);
            if (heavy[v] >= 0) self(self, heavy[v], true);
            for (int c : adj[v]) if (c != par[v] && c != heavy[v]) for (int i = tin[c]; i < tout[c]; i++) add(ver[i]);
            add(v); answer(v);                               // the structure holds exactly subtree(v)
            if (!keep) for (int i = tin[v]; i < tout[v]; i++) rem(ver[i]);
        };
        dfs(dfs, root, true);
    }
};
// example: for every subtree, the sum of the colours that occur most often (CF "Lomsat gelral")
vector<ll> dominatingColourSums(const vector<vector<int>>& adj, const vector<int>& col, int root) {
    int n = sz(adj); DsuOnTree D(adj, root); vector<int> cnt(n + 1); vector<ll> sumAt(n + 2), res(n); int best = 0;
    D.run(root,
        [&](int v) { int& c = cnt[col[v]]; sumAt[c] -= col[v]; c++; sumAt[c] += col[v]; best = max(best, c); },
        [&](int v) { int& c = cnt[col[v]]; sumAt[c] -= col[v]; c--; sumAt[c] += col[v]; if (!sumAt[best]) best--; },
        [&](int v) { res[v] = sumAt[best]; });
    return res;
}

Recursion depth equals the tree height; with $n=2\cdot 10^5$ path-shaped trees it is fine on most judges, otherwise raise the stack or emulate the recursion.

Proof: what the structure holds, and $O(n\log n)$

Invariant: every call dfs(v, keep) starts with an empty structure and returns with exactly subtree($v$) in it if keep, and empty otherwise. Induction: the light children are called first, each on an empty structure, and each leaves it empty; the heavy child is then called on an empty structure and leaves subtree(heavy); re-inserting the light subtrees and $v$ gives exactly subtree($v$), which answer(v) sees; a non-kept call then removes all of it. The root call starts empty. Cost: a vertex $u$ is added once when it is processed as $v$, and once more for every ancestor $v$ of which it lies in a light child’s subtree (the re-insertion loop); removals are matched by insertions. The light edges on $u$‘s root path number $\le {\log }_{2}n$ (each at least doubles the subtree size), so the total is $O(n\log n)$ operations.

Proof: most-frequent bookkeeping

sumAt[k] is the sum of colours with count exactly $k$, and best the maximum count. An add raises one count by one, so best can only increase by one, to that count. A remove lowers one count; if that colour was the only one at best, the new maximum is best - 1 (the same colour now has it). The order of operations is valid because removals happen only when an entire subtree is cleared, but correctness of best only needs the two local facts above.

Merging containers small into large

template <class T, class F> void smallToLarge(const vector<vector<int>>& adj, int root, vector<T> box, F answer) {
    // box[v]: v's own container; children are merged into the largest; answer(v, container) after merging
    int n = sz(adj); vector<int> par(n, -1), order{root}, id(n); iota(all(id), 0);   // id[v]: container currently owned by v
    for (int i = 0; i < n; i++) for (int v : adj[order[i]]) if (v != par[order[i]]) par[v] = order[i], order.push_back(v);
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        for (int c : adj[v]) if (c != par[v]) {
            if (box[id[v]].size() < box[id[c]].size()) swap(id[v], id[c]);
            for (auto& x : box[id[c]]) box[id[v]].insert(x);
            box[id[c]].clear();
        }
        answer(v, box[id[v]]);
    }
}

Use it with set<int> for “distinct values per subtree”, map<int, int> for frequencies, or a heap for “best $k$ values in a subtree”. The same principle drives union-by-size DSU, merging sorted lists in offline divide and conquer, and “merge components, relabel the smaller one”.

Proof: each element moves $O(\log n)$ times

When an element moves from container $B$ into $A$, $|A|\ge |B|$ before the merge, so the container now holding it has at least twice the size of the one it left. Sizes are at most $n$, so an element moves at most ${\log }_{2}n$ times; with $O(\log n)$ per insertion into an ordered set this is $O(n{\log }^{2}n)$ overall (with a hash set, $O(n\log n)$ expected). Answers are correct because after processing $v$, its container is the union of its own data and all children’s containers, i.e. its subtree.