树上启发式合并(dsu on tree)

upd: 文尾更新了一个更好的讲解, 但英文

当要维护树中所有以每个节点为根的子树的某个值, 而且以 rot 为根的树的值可以从以 rot 的 son 为根的子树转移过来.

eg => U41492 树上数颜色.

此时如果每往上传递子树的贡献时, 都暴力枚举子树中所有点, 复杂度就到了 $n^2$.

不妨这样做: 每往上传递子树的贡献时, 其他 ( 轻子树 )的仍旧暴力枚举所有点来计算贡献. 最大的那个子树 ( 重子树) 留到最后一个跑 dfs, 此时跑出来的结果不再动他, 直接往上回溯, 就传承给了上一级. 形象地说, 就是所有轻子树合并到了重子树上, 作为 fa 的子树的贡献.

可以证明这样的复杂度是 $nlogn$. 见代码.

https://www.luogu.com.cn/problem/U41492
void dfs(int x, int fa) {
    for (auto c : eg[x]) {
        if (c == fa)continue;
        if (c == son[x])continue;
        dfs(c, x); // dfs 跑完所有轻子树内部的答案
        clear(c, x); //注意: 跑完后我们要清空这个轻子树的贡献, 不能影响计算其他轻子树内的贡献.
    }
    if (son[x]) // 计算重子树的答案, 不去clear, 直接保留, 这样就传给了上一层.
        dfs(son[x], x);

    for (auto c : eg[x]) { // 暴力计算所有轻子树的贡献
        if (c == fa)continue;
        if (c == son[x])continue;
        add(c, x);
    }
    if (!num[clo[x]])sum++; // 根节点本身
    num[clo[x]]++;
    ans[x] = sum;
}

就是说: 遍历所有子树的贡献时, 由于换一根子树 dfs 之前就要清除之前子树的贡献, 我们可以选择一个子树作为最后跑 dfs 的, 这样他就是不用清除的, 就可以节省复杂度.

全部代码:

//AC https://www.luogu.com.cn/problem/U41492
#define int long long
const int Maxn = 1e5 + 10;

vector<int>eg[Maxn];
int n, v, u, sum = 0;
int clo[Maxn], num[Maxn], son[Maxn], ans[Maxn];

int get_son(int x, int fa) {
    int Max = 0, pos = 0, sum = 1;
    for (auto c : eg[x]) {
        if (c == fa)continue;
        int tp = get_son(c, x);
        sum += tp;
        if (tp > Max)Max = tp, pos = c;
    }
    if (pos)son[x] = pos;
    return sum;
}
void clear(int x, int fa) {
    num[clo[x]]--;
    if (!num[clo[x]])sum--;
    for (auto c : eg[x]) {
        if (c == fa)continue;
        clear(c, x);
    }
}
void add(int x, int fa) {
    if (!num[clo[x]])sum++;
    num[clo[x]]++;
    for (auto c : eg[x]) {
        if (c == fa)continue;
        add(c, x);
    }
}
void dfs(int x, int fa) {
    for (auto c : eg[x]) {
        if (c == fa)continue;
        if (c == son[x])continue;
        dfs(c, x);
        clear(c, x);
    }
    if (son[x])
        dfs(son[x], x);

    for (auto c : eg[x]) {
        if (c == fa)continue;
        if (c == son[x])continue;
        add(c, x);
    }
    if (!num[clo[x]])sum++;
    num[clo[x]]++;
    ans[x] = sum;
}
signed main()
{
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    cin >> n;
    for (int i = 1; i < n; ++i) {
        cin >> v >> u;
        eg[v].emplace_back(u);
        eg[u].emplace_back(v);
    }
    for (int i = 1; i <= n; ++i) {
        cin >> clo[i];
    }

    get_son(1, 1);
    dfs(1, 1);

    int m; cin >> m;
    for (int i = 1; i <= m; ++i) {
        cin >> v;
        cout << ans[v] << endl;
    }
}

另一题: E. Lomsat gelral, 直接看代码:

vector<int>edge[N];//邻接表存边
int son[N], sz[N]; //每个点的重儿子   每个点子树大小
int c[N];//每个节点的颜色
int num[N];//每种颜色的数量
ll ans[N];//存每个点答案
ll sum = 0; //存当前对应答案
int maxn = 0; //存对应最大值
void get_son(int x, int pre) { //预处理获取每个点的重儿子
    sz[x] = 1;
    for (int i = 0; i < edge[x].size(); i++) {
        int tmp = edge[x][i];
        if (tmp == pre) {
            continue;
        }
        get_son(tmp, x);
        sz[x] += sz[tmp];
        if (sz[tmp] > sz[son[x]]) { //子树最大的为重儿子
            son[x] = tmp;
        }
    }
}
void add(int x, int pre, int dep) { //更新除重儿子外的点（轻儿子+父亲）
    num[c[x]]++;//更新
    if (num[c[x]] > maxn) { //最大值被更新了
        maxn = num[c[x]];
        sum = c[x];
    }
    else if (num[c[x]] == maxn) { //出现相同的最大值
        sum += c[x];
    }
    for (int i = 0; i < edge[x].size(); i++) {
        int tmp = edge[x][i];
        if (tmp == pre || dep == 1 && tmp == son[x]) { //第一层的重儿子在dfs中更新过且未被清空，所以这里跳过
            continue;
        }
        add(tmp, x, dep + 1);
    }
}
void clear(int x, int pre) {
    num[c[x]]--;
    for (int i = 0; i < edge[x].size(); i++) {
        int tmp = edge[x][i];
        if (tmp == pre) {
            continue;
        }
        clear(tmp, x);
    }
}
void dfs(int x, int pre) { //当前dfs处理以x为根的答案
    for (int i = 0; i < edge[x].size(); i++) {
        int tmp = edge[x][i];
        if (tmp == pre || tmp == son[x]) { //先跳过重儿子，留到最后处理
            continue;
        }
        dfs(tmp, x); //处理x的轻儿子
        clear(tmp, x); //将处理x轻儿子中标记过的点清空防止干扰其他儿子的求解
        maxn = sum = 0; //有清空操作，要把这两个重置
    }
    if (son[x]) {
        dfs(son[x], x); //重儿子最后判断，不需要清空了
    }
    add(x, pre, 1); //将除重儿子外的点（轻儿子+父亲）重新更新进来
    ans[x] = sum;
}
int main() {
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0); //同步流
    int n;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
    }
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    get_son(1, 0); //预处理
    dfs(1, 0); //dfs搜答案
    for (int i = 1; i <= n; i++) {
        cout << ans[i] << " ";
    }
    return 0;
}

最后分享一个帖子, 里面有更多好讲解,好题
[Explanation] dsu on trees (small to large)