题目
lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义s(i,j) 为i 到j 的颜色数量。以及
现在他想让你求出所有的sum[i]
输入格式
第一行为一个整数n，表示树节点的数量
第二行为n个整数，分别表示n个节点的颜色c[1],c[2]……c[n]
接下来n-1行，每行为两个整数x,y，表示x和y之间有一条边
输出格式
输出n行，第i行为sum[i]
输入样例
5
1 2 3 2 3
1 2
2 3
2 4
1 5
输出样例
10
9
11
9
12
提示
sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10
sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9
sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11
sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9
sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12
对于40%的数据，n<=2000
对于100%的数据，1<=n,c[i]<=10^5
题解
明显点分治即可
对于每棵分治出来的树，考虑过根的所有路径对树内点的影响
首先单独考虑一种颜色的影响，从根节点出发到每棵子树的每个点(u)，(u)节点在该颜色下会产生贡献当且仅当(u)到根的路径上有该颜色的节点
所以我们只要找出一个子树中所有颜色为该颜色，且其祖先中没有该颜色【也就是最高的该颜色点】，其子树所有点都会产生贡献，那么所有的对根的贡献就是所有这样点的子树大小之和
考虑对子树内的点，就减去该子树的贡献，就转化为和根类似的了
每当第一次经过一种颜色的点时，其子树内所有点经过该点必定产生该颜色的贡献，此时把该颜色的贡献改为剩余子树的大小即可
还有，根节点的颜色特殊考虑
不知讲清楚没有，仔细想想还是很明显的
不过写起来细节真的多
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LL long long int
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define REP(i,n) for (int i = 1; i <= (n); i++)
using namespace std;
const int maxn = 100005,maxm = 200005,INF = 1000000000;
inline int read(){
int out = 0,flag = 1; char c = getchar();
while (c < 48 || c > 57){if (c == '-') flag = -1; c = getchar();}
while (c >= 48 && c <= 57){out = (out << 3) + (out << 1) + c - 48; c = getchar();}
return out * flag;
}
int h[maxn],ne = 2;
struct EDGE{int to,nxt;}ed[maxm];
inline void build(int u,int v){
ed[ne] = (EDGE){v,h[u]}; h[u] = ne++;
ed[ne] = (EDGE){u,h[v]}; h[v] = ne++;
}
int n,c[maxn],F[maxn],N,Siz[maxn],rt,vis[maxn],fa[maxn];
int id[maxn],st[maxn],top,Vis[maxn],now,tot,tots;
LL D[maxn],Sum[maxn],sum[maxn],Sumt[maxn],ttt;
void getrt(int u){
Siz[u] = 1; F[u] = 0;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]){
fa[to] = u; getrt(to);
Siz[u] += Siz[to];
F[u] = max(F[u],Siz[to]);
}
F[u] = max(F[u],N - Siz[u]);
if (F[u] < F[rt]) rt = u;
}
void dfs1(int u){
Siz[u] = 1;
if (!id[c[u]]) st[++top] = c[u],id[c[u]] = top;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]){
fa[to] = u; dfs1(to);
Siz[u] += Siz[to];
}
}
void dfs2(int u){
int p = id[c[u]],flag = 0;
if (p != 1 && Vis[p] != now) Sum[p] += Siz[u],Vis[p] = now,flag = 1;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs2(to);
if (flag) Vis[p] = 0;
}
void dfs3(int u){
int p = id[c[u]],flag = 0;
if (p != 1 && Vis[p] != now) Sumt[p] += Siz[u],ttt += Siz[u],Vis[p] = now,flag = 1;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs3(to);
if (flag) Vis[p] = 0;
}
void dfs4(int u){
int p = id[c[u]],flag = 0;
D[u] = D[fa[u]];
if (p != 1 && Vis[p] != now){
D[u] -= (Sum[p] - Sumt[p]) - (tot - tots);
Vis[p] = now; flag = 1;
}
sum[u] += D[u] + (tot - tots);
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs4(to);
if (flag) Vis[p] = 0;
}
void dfs5(int u){
Sumt[id[c[u]]] = 0;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs5(to);
}
void dfs6(int u){
D[u] = 0;
Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs6(to);
}
void solve(int u){
vis[u] = true; Siz[u] = 1;
st[top = 1] = c[u]; id[c[u]] = top;
Redge(u) if (!vis[to = ed[k].to]){
fa[to] = u; dfs1(to);
Siz[u] += Siz[to];
}
now = 0; tot = Siz[u];
Redge(u) if (!vis[to = ed[k].to]){
now++; dfs2(to);
}
REP(i,top) D[u] += Sum[i];
sum[u] += D[u] + Siz[u];
Redge(u) if (!vis[to = ed[k].to]){
now++; ttt = 0; tots = Siz[to]; dfs3(to);
now++; D[u] -= ttt; dfs4(to);
D[u] += ttt; now++; dfs5(to);
}
D[u] = 0;
Redge(u) if (!vis[to = ed[k].to]) dfs6(to);
REP(i,top) Vis[i] = Sum[i] = Sumt[i] = id[st[i]] = 0;
Redge(u) if (!vis[to = ed[k].to]){
N = Siz[to]; F[rt = 0] = INF;
getrt(to); solve(rt);
}
}
int main(){
n = read();
REP(i,n) c[i] = read();
for (int i = 1; i < n; i++) build(read(),read());
F[rt = 0] = INF; N = n;
getrt(1); solve(rt);
REP(i,n) printf("%lld
",sum[i]);
return 0;
}