难度：NOI/NOI+/CTSC

题目描述

        lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义$s(i,j)$ 为i 到j 的颜色数量。以及

        现在他想让你求出所有的sum[i]。

输入输出格式

输入格式

        第一行为一个整数n，表示树节点的数量，$n,c[i]\leq 10^5$
        第二行为n个整数，分别表示n个节点的颜色$c[1],c[2]……c[n]$
        接下来n-1行，每行为两个整数x,y，表示$x$和$y$之间有一条边

输出格式

        输出n行，第i行为sum[i]

输入输出样例

Sample input

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

Sample output

10
9
11
9
12

题解

        有难度的点分治，强烈推荐一做。
        首先先保证复杂度。对于每一次分治，我们必须在$O(n)$复杂度下完成处理。假设现在分治到以rt为根的子树，下面需要在$O(n)$复杂度下求所有经过rt点的路径对答案的贡献。
        点分治的实质是将所有路径划分为两类：一类经过根结点，一类不经过，然后用分治的方法统计两类路径对答案的贡献（还记得吗？）。那么如何统计答案？
        首先认识一个事实：如果某一个点是其到根结点（不含根结点）的路径上第一次出现某种点权值的点，那么这个点将对其它点贡献答案size，这里的size是其子树上的点的数量。
        我们用一遍DFS统计每一种权值的对应贡献，这个贡献用一个数组存起来（设为color，color[x]表示权值x对应的答案贡献）。然后求它们的和，记为sum。
        下面遍历每一个子树。对于某一棵子树，对其进行DFS，删掉这个树上所有点对上面求出的数组以及sum的贡献。经过这一步操作得到的贡献数组以及sum就不再包含这个子树了。做这一步的目的是，使两个点对不出现在同一棵子树（同一棵子树上的点对路径不经过根，不需要统计）。
        对于这一棵子树，再进行一步遍历，记录所经过点到根（不含根）的所有点权种类，然后我们可以这样统计答案：

        这里求和所用的$i$是出现在路径上的权值，$p$代表出现在路径上的权值种类数，$n$是整棵树（正在处理的树）的总点数，$size$是当前子树的点数。这个公式是易于理解的。
        要注意的是，为了方便起见，在整个过程中，我们都不统计根结点权值对答案的贡献，因此当遍历到某一个权值与根相同的点时，认为它无点权，不计入贡献，也不统计种类数量。对于所有经过根的路径，根结点点权必然会计入答案，这样的话，我们只需要在答案贡献时统一加上一就好了。
        最后需要单独统计根结点的答案。这个实现起来很简单，其贡献的答案就是$sum+size$。理解：其它权值的贡献加上自己本身权值种类的贡献。
        在分治过程中需要不断清数组，为了保证复杂度，不要用memset清空，这是40分与100分AC的差别。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#include <bits/stdc++.h>

#define N 100005
using namespace std;
struct Edge {
    int next, to;
} edge[N << 1];
int head[N], n, c[N], cnt = 1, vis[N], rt, minn, root_size[N], np, hs[N], numOfColor, objColor, size[N];
int tempRoot;
typedef long long ll;
ll ans[N], color[N], sum, now;

inline void add(int x, int y) {
    edge[cnt].to = y, edge[cnt].next = head[x], head[x] = cnt++;
}

void DFS_root(int x, int fa) {
    root_size[x] = 1;
    int ff = 0;
    for (int i = head[x]; i; i = edge[i].next) {
        if (!vis[edge[i].to] && edge[i].to != fa) {
            DFS_root(edge[i].to, x), root_size[x] += root_size[edge[i].to];
            ff = max(ff, root_size[edge[i].to]);
        }
    }
    ff = max(ff, np - root_size[x]);
    if (ff < minn)minn = ff, rt = x;
}

inline void findRoot(int x, int num) {
    np = num, minn = 0x7fffffff, DFS_root(x, 0);
}

void DFS(int x, int fa, int type) {
    size[x] = 1;
    if (c[x] != objColor)++hs[c[x]];
    for (int i = head[x]; i; i = edge[i].next) {
        if (!vis[edge[i].to] && edge[i].to != fa)DFS(edge[i].to, x, type), size[x] += size[edge[i].to];
    }
    if (c[x] != objColor && hs[c[x]] == 1)color[c[x]] += type * size[x], sum += type * size[x];
    if (c[x] != objColor)--hs[c[x]];
}

void DFS_findAns(int x, int fa) {
    if (c[x] != objColor)++hs[c[x]];
    if (c[x] != objColor && hs[c[x]] == 1)++numOfColor, now += color[c[x]];
    ans[x] += sum - now + (numOfColor + 1) * (size[rt] - size[tempRoot]);//手动加一，计入根结点贡献
    for (int i = head[x]; i; i = edge[i].next)if (!vis[edge[i].to] && edge[i].to != fa)DFS_findAns(edge[i].to, x);
    if (c[x] != objColor && hs[c[x]] == 1)--numOfColor, now -= color[c[x]];
    if (c[x] != objColor)--hs[c[x]];
}

void DFS_clear(int x, int fa) {
    if (c[x] != objColor)color[c[x]] = 0;
    for (int i = head[x]; i; i = edge[i].next) {
        if (edge[i].to != fa && !vis[edge[i].to])DFS_clear(edge[i].to, x);
    }
}

void divide(int x) {//点分治主函数
    objColor = c[x], vis[x] = 1, sum = 0, DFS(x, 0, 1);//objColor是根的权值，不计入答案。第一遍DFS求sum和color数组
    for (int i = head[x]; i; i = edge[i].next) {
        if (!vis[edge[i].to]) {//记录临时根（就是子树的根），一遍DFS清除子树贡献
            tempRoot = edge[i].to, DFS(edge[i].to, x, -1), numOfColor = 0, now = 0;
            DFS_findAns(edge[i].to, x), DFS(edge[i].to, x, 1);//另一遍DFS找答案，最后一遍DFS恢复数据（再加回来）
        }
    }
    ans[x] += sum + size[x], DFS_clear(x, 0);//统计根结点答案，用DFS清空数组
    for (int i = head[x]; i; i = edge[i].next)if (!vis[edge[i].to])findRoot(edge[i].to, size[edge[i].to]), divide(rt);
}

int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++)scanf("%d", c + i);
    for (int i = 1, x, y; i < n; i++) scanf("%d%d", &x, &y), add(x, y), add(y, x);
    findRoot(1, n), divide(rt);
    for (int i = 1; i <= n; i++)printf("%lld\n", ans[i]);
    return 0;
}