← all categories

Small To Large 108ae9a3

Small-to-large merging ("DSU on tree"): merge a child's set into the larger of the two and re-insert the smaller one's elements, giving O(n log^2 n) total instead of O(n^2) for per-node subtree aggregate queries. Worked example: count of distinct colors in each node's subtree, given colors a[] and adjacency adj (rooted at 0).

Time: O(n \log^2 n). tested

content/graphs/small-to-large.h

vector<int> distinctColorCounts(int n, vector<int> &a, vector<vector<int>> &adj) {
    vector<set<int>> col(n);
    vector<int> ans(n);
    auto dfs = [&](auto &&dfs, int u, int p) -> void {
        for (int v : adj[u]) {
            if (v == p) continue;
            dfs(dfs, v, u);
            if (col[u].size() < col[v].size()) swap(col[u], col[v]);
            for (int x : col[v]) col[u].insert(x);
        }
        col[u].insert(a[u]);
        ans[u] = (int)col[u].size();
    };
    dfs(dfs, 0, -1);
    return ans;
}