大致题意： 给定一棵树，每个节点有一个颜色，定义\(s(i,j)\)为\(i\)到\(j\)路径上颜色数量，请你对于每一个\(i\)求出\(\sum_{j=1}^n s(i,j)\)。
点分治
这种题目比较显然是点分治吧... ...
大致思路
首先，按照点分治的基本套路，对于一棵子树内的路径，我们分两种情况讨论：
- 经过根节点的
- 不经过根节点的
呃，第二种情况就按照点分治继续处理下去就可以了，因此忽略不提。
而对于第一种情况的边，我们又可以分为两类：
- 以根节点为一个端点的
- 不以根节点为端点的
接下来就按此分类处理即可。
先是一波定义
首先，我们需要先来一波定义。
\(sum\)：所有颜色造成的总贡献。
\(col_i\)：节点\(i\)的颜色。
\(Val_i\)：颜色\(i\)对答案造成的贡献。
\(Size_i\)：以节点\(i\)为根的子树大小。
具体流程
-
首先是一遍\(dfs\)预处理。在这一遍与处理中，我们主要求出两个东西：\(Size\)和\(Val\)两个数组，当然，在求\(Val\)数组的同时，也要顺带求出\(sum\)。
-
然后，我们就可以对第一类路径进行处理，更新\(ans_{rt}\)。
-
具体操作：\(ans_{rt}+=sum-Val_{col_{rt}}+Size_{rt}\)。
-
这样做的理由是先将含有\(rt\)的颜色所造成的贡献从\(sum\)中删去（不然会造成重复计算），然后加上\(Size_{rt}\)（即含有\(rt\)的路径条数，\(rt\)的颜色在这些路径中会各对答案造成一点贡献）。
-
-
接下来便是对第二类路径的处理。我们考虑对于根节点的每一棵子树如何计算其贡献值。
- 第一步自然是将这棵子树中的贡献清空，不然会造成重复计算。
- 接下来，从该子节点开始，一直往下遍历，更新每一个节点的\(ans\)。
- 首先，我们要将沿途经过的在该子树中出现的颜色的贡献从\(sum\)中减去，避免重复计算，并用一个变量\(color\_tot\)记录该子树中颜色种数。
- 则对于当前节点\(x\)，我们可以将\(ans_x\)加上\(sum+color\_tot*Other\)。
- 其中\(Other\)表示除该棵子树外其他子树以及根节点的总点数，因此将剩余贡献加上\(color\_tot*Other\)就是当前节点能造成的贡献了（和对第一类路径的操作同理）。
-
好吧，还有一个很重要的细节：注意清空数组。
- 但是，如何清空数组也是有学问的。
- 如果直接\(memset\)，显然复杂度成了\(O(n)\)，可能会\(TLE\)。
- 所以，我们直接将该子树中出现的所有颜色的相关值全部清空即可，就是一个遍历子树的过程，时间复杂度为\(O(子树大小)\)，从而保证了时间复杂度。
代码
#include<bits/stdc++.h>
#define max(x,y) ((x)>(y)?(x):(y))
#define min(x,y) ((x)<(y)?(x):(y))
#define uint unsigned int
#define LL long long
#define ull unsigned long long
#define swap(x,y) (x^=y,y^=x,x^=y)
#define abs(x) ((x)<0?-(x):(x))
#define INF 1e9
#define Inc(x,y) ((x+=(y))>=MOD&&(x-=MOD))
#define ten(x) (((x)<<3)+((x)<<1))
#define N 100000
#define add(x,y) (e[++ee].nxt=lnk[x],e[lnk[x]=ee].to=y)
using namespace std;
LL n,ee=0,lnk[N+5],col[N+5];
struct edge
{
LL to,nxt;
}e[(N<<1)+5];
class FIO
{
private:
#define Fsize 100000
#define tc() (FinNow==FinEnd&&(FinEnd=(FinNow=Fin)+fread(Fin,1,Fsize,stdin),FinNow==FinEnd)?EOF:*FinNow++)
#define pc(ch) (FoutSize<Fsize?Fout[FoutSize++]=ch:(fwrite(Fout,1,FoutSize,stdout),Fout[(FoutSize=0)++]=ch))
LL f,FoutSize,OutputTop;char ch,Fin[Fsize],*FinNow,*FinEnd,Fout[Fsize],OutputStack[Fsize];
public:
FIO() {FinNow=FinEnd=Fin;}
inline void read(LL &x) {x=0,f=1;while(!isdigit(ch=tc())) f=ch^'-'?1:-1;while(x=ten(x)+(ch&15),isdigit(ch=tc()));x*=f;}
inline void read_char(char &x) {while(isspace(x=tc()));}
inline void read_string(string &x) {x="";while(isspace(ch=tc()));while(x+=ch,!isspace(ch=tc())) if(!~ch) return;}
inline void write(LL x) {if(!x) return (void)pc('0');if(x<0) pc('-'),x=-x;while(x) OutputStack[++OutputTop]=x%10+48,x/=10;while(OutputTop) pc(OutputStack[OutputTop]),--OutputTop;}
inline void write_char(char x) {pc(x);}
inline void write_string(string x) {register LL i,len=x.length();for(i=0;i<len;++i) pc(x[i]);}
inline void end() {fwrite(Fout,1,FoutSize,stdout);}
}F;
class Class_DotSolver//点分治
{
private:
LL rt,sum,color_tot,used[N+5],Size[N+5],Max[N+5],cnt[N+5],Val[N+5];
inline void GetRt(LL x,LL lst,LL tot)//找树的重心，确定根节点
{
register LL i;
for(i=lnk[x],Size[x]=1,Max[x]=0;i;i=e[i].nxt)
if(e[i].to^lst&&!used[e[i].to]) GetRt(e[i].to,x,tot),Size[x]+=Size[e[i].to],Max[x]=max(Max[x],Size[e[i].to]);
if((Max[x]=max(Max[x],tot-Size[x]))<Max[rt]) rt=x;
}
inline void Clear(LL x,LL lst)//清空数组
{
for(register LL i=lnk[x];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Clear(e[i].to,x);//枚举子节点，遍历该子树
cnt[col[x]]=Val[col[x]]=0;//清空与当前节点颜色相关的数据
}
inline void Init(LL x,LL lst)//初始化出子树大小和贡献
{
register LL i;
for(i=lnk[x],Size[x]=1,++cnt[col[x]];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Init(e[i].to,x),Size[x]+=Size[e[i].to];//枚举子节点进行处理
if(!--cnt[col[x]]) sum+=Size[x],Val[col[x]]+=Size[x];//如果该颜色是第一次出现，将当前颜色造成的贡献加上该子树大小
}
inline void Change(LL x,LL lst,LL flag)//删除（flag=0）or 恢复（flag=1）一棵子树中颜色造成的贡献
{
register LL i;
for(i=lnk[x],++cnt[col[x]];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Change(e[i].to,x,flag);
if(!--cnt[col[x]]) flag?(sum+=Size[x],Val[col[x]]+=Size[x]):(sum-=Size[x],Val[col[x]]-=Size[x]);//与上面的操作同理
}
inline void F5(LL x,LL lst,LL Other)//更新一棵子树中的ans
{
register LL i;
if(!cnt[col[x]]++) sum-=Val[col[x]],++color_tot;//如果是第一次出现，将当前颜色的贡献减去，然后将该子树中颜色种数加1
for(i=lnk[x],ans[x]+=sum+color_tot*Other;i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) F5(e[i].to,x,Other);
if(!--cnt[col[x]]) sum+=Val[col[x]],--color_tot;//撤销操作
}
inline void Solve(LL x)//处理以x为根的子树
{
register LL i,j;
for(i=lnk[x],used[x]=1,Init(x,0),ans[x]+=sum-Val[col[x]]+Size[x];i;i=e[i].nxt)//初始化+更新根节点的ans
{
if(used[e[i].to]) continue;
++cnt[col[x]],sum-=Size[e[i].to],Val[col[x]]-=Size[e[i].to],Change(e[i].to,x,0),//先清空该子树造成的贡献
--cnt[col[x]],F5(e[i].to,x,Size[x]-Size[e[i].to]),++cnt[col[x]],//更新该子树内ans
Change(e[i].to,x,1),Val[col[x]]+=Size[e[i].to],sum+=Size[e[i].to],--cnt[col[x]];//恢复该子树造成的贡献
}
for(i=lnk[x],sum=color_tot=0,Clear(x,0);i;i=e[i].nxt)//清空
if(!used[e[i].to]) GetRt(e[i].to,rt=0,Size[e[i].to]),Solve(rt);//继续处理子树信息
}
public:
Class_DotSolver() {Max[0]=INF;}
LL ans[N+5];
inline void GetAns() {GetRt(1,rt=0,n),Solve(rt);}
}DotSolver;
int main()
{
register LL i,x,y;
for(F.read(n),i=1;i<=n;++i) F.read(col[i]);
for(i=1;i<n;++i) F.read(x),F.read(y),add(x,y),add(y,x);
for(DotSolver.GetAns(),i=1;i<=n;++i) F.write(DotSolver.ans[i]),F.write_char('\n');
return F.end(),0;
}