题目描述

难度分:$2470$

输入$n(1 \leq n \leq 2 \times 10^5)$和长为$n$的数组$a(1 \leq a[i] \leq n)$。

然后输入一棵无向树的$n-1$条边，节点编号从$1$到$n$。节点$i$的颜色是$a[i]$。

定义$f(c)=$包含颜色$c$的简单路径的数目。注：只有$1$个点也算路径。

输出$f(1),f(2),…,f(n)$。

输入样例$1$

3
1 2 1
1 2
2 3

输出样例$1$

5
4
0

输入样例$2$

1
1

输出样例$2$

1

输入样例$3$

2
1 2
1 2

输出样例$3$

2
2

输入样例$4$

5
1 2 3 4 5
1 2
2 3
3 4
3 5

输出样例$4$

5
8
10
5
5

输入样例$5$

8
2 7 2 5 4 1 7 5
3 1
1 2
2 7
4 5
5 6
6 8
7 8

输出样例$5$

18
15
0
14
23
0
23
0

算法

逆向思维

正难则反，计算不包含颜色$c$的简单路径数。所有路径数减去不包含颜色$c$的路径数就是答案。

去掉颜色$c$节点后，树分成了若干连通块。对于大小为$m$的连通块，其中有$\frac{m \times (m+1)}{2}$条简单路径。接下来的问题就在于如何快速计算各个连通块的大小？

如图，去掉粉色节点后，考虑包含节点$y$的连通块，它的大小等于：树$y$的大小，减去子树$z_1$，$z_2$，$z_3$的大小之和。

如何计算子树$z_1$，$z_2$，$z_3$的大小之和？额外用一个数组$size[c]$记录以颜色$c$为根的子树的大小之和。

但这样还有问题，考虑$x$更上面的节点，这样做$z_1$，$z_2$，$z_3$的子树大小就会和$x$的子树大小相加。为了避免重复累加子树大小，直接用子树$x$的大小覆盖子树$z_1$，$z_2$，$z_3$的大小之和。

在实现上，于$dfs(x)$的开头用临时变量$old$记录$size[c]$；在$dfs(x)$的末尾，覆盖$size[c]=old+$子树$x$的大小。

复杂度分析

时间复杂度

DFS遍历整棵树就能预处理出计算答案所需要的信息，接下来遍历$i \in [1,n]$，对每个$i$节点$O(1)$计算答案即可。因此，整个算法的时间复杂度为$O(n)$。

空间复杂度

$size$数组和$ans$数组都是线性空间消耗；整棵树的邻接表空间消耗也是$O(n)$；在DFS的过程中，如果整棵树退化成链，递归深度就会达到$O(n)$，空间复杂度仍然是$O(n)$。因此，整个算法的额外空间复杂度为$O(n)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;
const int N = 200010;
int n, c[N];
long long ans[N], Size[N];
vector<int> graph[N];

int dfs(int u, int fa) {
    int old = Size[c[u]];
    int sz = 1;
    for(int v: graph[u]) {
        if(v == fa) continue;
        Size[c[u]] = 0;
        int szW = dfs(v, u);
        sz += szW;
        int m = szW - Size[c[u]];
        ans[c[u]] += m*(m + 1LL)/2;
    }
    Size[c[u]] = old + sz;
    return sz;
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &c[i]);
        graph[i].clear();
        Size[i] = ans[i] = 0;
    }
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; i++) {
        int m = n - Size[i];
        printf("%lld\n", n*(n + 1LL)/2 - ans[i] - m*(m + 1LL)/2);
    }
    return 0;
}