T176273 Hazardous【dsu on tree】

最新推荐文章于 2022-05-12 19:43:15 发布

_7许

最新推荐文章于 2022-05-12 19:43:15 发布

阅读量190

点赞数

分类专栏：图论 # dsu on tree

本文链接：https://blog.csdn.net/weixin_44049850/article/details/116208598

版权

图论同时被 2 个专栏收录

53 篇文章 0 订阅

订阅专栏

dsu on tree

3 篇文章 0 订阅

订阅专栏

题目链接

题意：

给定一棵 $n$ 个结点的树，根固定为 $1$ ，每个结点有一个颜色。
对于结点 $u$ ，它的值定义为子树 $u$ 中任意两个相同颜色的结点的距离之和。我们需要得到每棵子树的值并输出。
在这里插入图片描述

思路：

对于子树 $u$ ，我们假设它有 $x$ 个儿子，我们假定这些儿子的编号为 $1, 2, 3, . . ., x$
我们假设每个结点 $i$ 的颜色是 $c [i]$ ，深度是 $d [i]$

我们统计一个后缀和 $s u m [i]$ ：表示目前为止颜色为 $i$ 的结点的深度和；统计颜色数量 $n u m [i]$ ：表示目前为止颜色为 $i$ 的结点的数量。

当然子树问题首先想到 $d s u o n t r e e$

对于子树 $u$ ，它每个儿子子树对于结点 $u$ 的贡献都是 $s u m [c [u]] - n u m [c [u]] * d [u]$ （当然这里的 $s u m []$ 和 $n u m []$ 只是每一棵子树的单独的贡献）这个贡献的话就可以通过先遍历轻子树（清除轻子树的贡献）再遍历重子树的方式得到

那么子树 $u$ 每个儿子子树之间产生的贡献怎么算呢？

我们先统计重子树的贡献，然后依次暴力统计轻子树的贡献，在这个过程中我们只需要不清除每棵子树对于 $s u m []$ 和 $n u m []$ 的贡献即可
每棵轻子树上的结点 $v$ 对于已统计过的子树的贡献是 $s u m [c [v]] - n u m [c [v]] * d [u] + (d [v] - d [u]) * n u m [c [v]]$

需要注意的是需要把每棵轻子树上的所有结点的贡献都统计完了之后再将该子树对于 $s u m []$ 和 $n u m []$ 的贡献统计上才行。

样例
5
1 2 1 1 2
1 2
2 3
2 4
5 1
 ans: 8 2 0 0 0
4
1 1 1 1
1 2
1 3
1 4
 ans: 9 0 0 0
5
1 1 1 2 3
1 2
2 3
2 4
4 5
 ans: 4 1 0 0 0
6
1 1 1 2 1 2
1 2
2 3
2 4
2 5
1 6
 ans: 12 4 0 0 0 0
11
1 1 1 2 1 2 1 2 3 1 3
1 2
2 3
2 4
2 5
1 6
1 8
8 9
8 10
8 11
9 7
 ans: 52 4 0 0 0 0 0 5 0 0 0

#include <bits/stdc++.h>
#define eps 1e-6
#define INF 0x3f3f3f3f3f3f3f3f
using namespace std;
typedef long long ll;

ll read() {
    ll x = 0, f = 1; char ch = getchar();
    while(ch < '0' || ch > '9') { if(ch == '-') f = -f; ch = getchar(); }
    while(ch >= '0' && ch <= '9') { x = x * 10 + ch - '0'; ch = getchar(); }
    return x * f;
}
const int maxN = 1e5 + 10;
int n, c[maxN];
struct EDGE{
    int adj, to;
}edge[maxN << 1];
int head[maxN], cnt;
void add_edge(int u, int v) {
    edge[cnt] = EDGE{head[u], v};
    head[u] = cnt ++ ;
}
ll d[maxN];
int son[maxN], siz[maxN], nowSon;
void dfs(int u, int fa) {
    d[u] = d[fa] + 1;
    siz[u] = 1;
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[son[u]] < siz[v]) son[u] = v;
    }
}
ll ans[maxN]; //每个子树的答案
ll sum[maxN]; //每个颜色的后缀和
ll num[maxN]; //每个颜色的数量
void add(int u) {
    sum[c[u]] += d[u];
    num[c[u]] ++;
}
void update(int u, int fa) {
    add(u);
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        update(v, u);
    }
}
void del(int u, int fa) {
    num[c[u]] --;
    sum[c[u]] -= d[u];
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        del(v, u);
    }
}
void cal(int u, int fa, int ff) {
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        ans[ff] += sum[c[v]] - d[ff] * num[c[v]] + (d[v] - d[ff]) * num[c[v]];
        cal(v, u, ff);
        if(u == ff) update(v, u);
    }
}
void dsu(int u, int fa, bool z) {
    for(int i = head[u]; ~i; i = edge[i].adj ) {
        int v = edge[i].to;
        if(v == fa || v == son[u]) continue;
        dsu(v, u, false);
    }
    if(son[u]) {
        dsu(son[u], u, true);
        nowSon = son[u];
    }
    cal(u, fa, u);
    if(!nowSon) update(u, fa);
    else sum[c[u]] += d[u], num[c[u]] ++;
    nowSon = 0;
    ans[fa] += sum[c[fa]] - d[fa] * num[c[fa]];
    if(!z) del(u, fa);
}
void get(int u, int fa) {
    for(int i = head[u]; ~i; i = edge[i].adj) {
        int v = edge[i].to;
        if(v == fa) continue;
        get(v, u);
        ans[u] += ans[v];
    }
}
int main() {
    n = read();
    for(int i = 1; i <= n; ++ i ) {
        c[i] = read();
        head[i] = -1;
    }
    for(int i = 1; i < n; ++ i ) {
        int u = read(), v = read();
        add_edge(u, v);
        add_edge(v, u);
    }
    dfs(1, 0);
    dsu(1, 0, true);
    get(1, 0);
    for(int i = 1; i <= n; ++ i ) {
        printf("%lld%c", ans[i], " \n"[i == n]);
    }
    return 0;
}