树上启发式合并

用处

擅长处理对树上每个节点 \(u\) 求关于 \(u\) 的子树的信息的题目。

分析

一个例子，树上每个点都有一个颜色，要求对树上每个点 \(u\) 求出它子树内不同的颜色个数。

考虑暴力怎么做，\(O(n)\) 对每个 \(u\) 子树进行枚举，总共 \(n\) 个点，就是 \(O(n^2)\) 的复杂度。

不妨换一个暴力，考虑对每一个节点维护一个桶 \(t_{u,c}\) 表示 \(u\) 的子树里颜色 \(c\) 的出现次数，每次让 \(u\) 合并儿子节点 \(v\) 的桶，时空复杂度 \(O(nV)\)（\(V\) 就是颜色的大小）。

发现这个暴力还是垃圾，时空都暴，没关系，可以优化空间，只维护一个数组 \(t_c\) 表示颜色 \(c\) 出现次数。

现在就考虑复用这个数组，统计完 \(u\) 的子树信息、得到答案后，把它对数组 \(t\) 的影响清空（不清空，就会对 \(u\) 的兄弟节点造成影响）。

代码差不多这样子。

int cnt; // 不同颜色个数
int t[N], ans[N], col[N];

void add( int c) {
    if (! t[c] ++) cnt ++;
}

void del( int c) {
    if (! -- t[c]) cnt --;
}

void work( int u, int fu, int op) {
    if (op == 1) add(col[u]);
    else del(col[u]);

    for ( auto v : G[u])
        if (v != fu) work(v, u, op);
}

void dfs( int u, int fu) {
    for ( auto v : G[u])
        if (v != fu) dfs(v, u);

    work(u, fu, 1);
    ans[u] = cnt;
    work(u, fu, -1);
}

空间复杂度 \(O(n)\)，时间仍然是 \(O(n^2)\)。

想到一个小优化，如果 \(v\) 是自己的父亲 \(u\) 的最后枚举到的儿子，那么就不用清空对数组 \(t\) 的影响了，因为自己子树的信息，父亲 \(u\) 也包含。

看似这个优化屁用没有，但如果说钦定这个最后被枚举到的 \(v\) 为 \(u\) 的重儿子呢？

总结一下，如果 \(v\) 是轻儿子，那么在统计完自己的信息、得到答案后，还需要枚举一遍自己的子树清空对数组 \(t\) 的影响。如果 \(v\) 是重儿子，就不用管了（注意：重儿子是最后被枚举到的点）。然后在处理点 \(u\) 的答案时，只用统计它的轻儿子的子树信息即可，因为 \(u\) 的重儿子的信息已经被统计过。

分析一下复杂度。

如果连接 \(v\) 和它的父亲 \(u\) 的边是一条轻边，就还要在枚举一遍 \(v\) 的子树，撤销对数组 \(t\) 的影响，且处理 \(u\) 时还要枚举一次 \(v\) 的子树。

存在一个经典的结论，一个节点到根节点的路径上只会有 \(O(\log n)\) 条轻边。

所以一个点只会被枚举到最多 \(O(\log n)\) 次。

此刻这个方法升华为龙，复杂度 \(O(n\log n)\)。

这就是算法的魅力之处，一个小小的改动，对于复杂度就是质的飞升。

int cnt; // 不同颜色个数
int t[N], ans[N], col[N];

void add( int c) {
    if (! t[c] ++) cnt ++;
}

void del( int c) {
    if (! -- t[c]) cnt --;
}

int fa[N], son[N], siz[N];

int tot;
int dfn[N], ed[N], rev[N];

void init( int u, int fu) {
    fa[u] = fu, siz[u] = 1;
    dfn[u] = ++ tot, rev[dfn[u]] = u;

    for ( auto v : G[u]) {
        if (v == fu) continue ;
        init(v, u);
        siz[u] += siz[v], son[u] = (siz[v] > siz[son[u]] ? v : son[u]);
    }

    ed[u] = tot;
}

void dfs( int u, int op) {
    for ( auto v : G[u])
        if (v != fa[u] && son[u] != v) dfs(v, 0);

    if (son[u]) dfs(son[u], 1); // 最后枚举重儿子
    
    add(col[u]);

    for ( auto v : G[u]) // 统计轻儿子的子树信息来得到自己的答案（重儿子的信息已经被统计）
        if (v != fa[u] && son[u] != v)
            for ( int i = dfn[v]; i <= ed[v]; i ++)
                add(col[rev[i]]);

    ans[u] = cnt;

    if (! op) // 如果自己是轻儿子，也就是被轻边连接，那么就撤销自己的影响
        for ( int i = dfn[u]; i <= ed[u]; i ++)
            del(col[rev[i]]);
}

一些细节

如果统计信息的容器不能支持删除该怎么办？比如李超线段树（它不能删除已经插入的直线）。

那么可以给它动态开点，然后删除时，直接把它的所有节点清除即可。

这为什么对呢？因为重儿子最后被枚举，就会导致从重儿子开始回溯时，只会走重边上去，相当于走了条重链，那么轻儿子的清空总是在重儿子前。

所以当要清空轻儿子的信息时，容器中其实只有这个轻儿子的信息，直接清除完也不会有影响。