大致题意: 给定一棵树,每个节点有一个颜色,定义 s ( i , j ) s(i,j) s(i,j)为 i i i到 j j j路径上颜色数量,请你对于每一个 i i i求出 ∑ i = 1 n s ( i , j ) \sum_{i=1}^n s(i,j) ∑i=1ns(i,j)。
点分治
这种题目比较显然是点分治吧… …
L i n k Link Link
点分治 详见博客 初学点分治
大致思路
首先,按照点分治的基本套路,对于一棵子树内的路径,我们分两种情况讨论:
- 经过根节点的
- 不经过根节点的
呃,第二种情况就按照点分治继续处理下去就可以了,因此忽略不提。
而对于第一种情况的边,我们又可以分为两类:
- 以根节点为一个端点的
- 不以根节点为端点的
接下来就按此分类处理即可。
先是一波定义
首先,我们需要先来一波定义。
s u m sum sum:所有颜色造成的总贡献。
c o l i col_i coli:节点 i i i的颜色。
V a l i Val_i Vali:颜色 i i i对答案造成的贡献。
S i z e i Size_i Sizei:以节点 i i i为根的子树大小。
具体流程
-
首先是一遍 d f s dfs dfs预处理。在这一遍与处理中,我们主要求出两个东西: S i z e Size Size和 V a l Val Val两个数组,当然,在求 V a l Val Val数组的同时,也要顺带求出 s u m sum sum。
-
然后,我们就可以对第一类路径进行处理,更新 a n s r t ans_{rt} ansrt。
-
具体操作: a n s r t + = s u m − V a l c o l r t + S i z e r t ans_{rt}+=sum-Val_{col_{rt}}+Size_{rt} ansrt+=sum−Valcolrt+Sizert。
-
这样做的理由是先将含有 r t rt rt的颜色所造成的贡献从 s u m sum sum中删去(不然会造成重复计算),然后加上 S i z e r t Size_{rt} Sizert(即含有 r t rt rt的路径条数, r t rt rt的颜色在这些路径中会各对答案造成一点贡献)。
-
-
接下来便是对第二类路径的处理。我们考虑对于根节点的每一棵子树如何计算其贡献值。
- 第一步自然是将这棵子树中的贡献清空,不然会造成重复计算。
- 接下来,从该子节点开始,一直往下遍历,更新每一个节点的 a n s ans ans。
- 首先,我们要将沿途经过的在该子树中出现的颜色的贡献从 s u m sum sum中减去,避免重复计算,并用一个变量 c o l o r _ t o t color\_tot color_tot记录该子树中颜色种数。
- 则对于当前节点 x x x,我们可以将 a n s x ans_x ansx加上 s u m + c o l o r _ t o t ∗ O t h e r sum+color\_tot*Other sum+color_tot∗Other。
- 其中 O t h e r Other Other表示除该棵子树外其他子树以及根节点的总点数,因此将剩余贡献加上 c o l o r _ t o t ∗ O t h e r color\_tot*Other color_tot∗Other就是当前节点能造成的贡献了(和对第一类路径的操作同理)。
-
好吧,还有一个很重要的细节:注意清空数组。
- 但是,如何清空数组也是有学问的。
- 如果直接 m e m s e t memset memset,显然复杂度成了 O ( n ) O(n) O(n),可能会 T L E TLE TLE。
- 所以,我们直接将该子树中出现的所有颜色的相关值全部清空即可,就是一个遍历子树的过程,时间复杂度为 O ( 子 树 大 小 ) O(子树大小) 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;
}