Submission #35854775

---

Source Code Expand

// LUOGU_RID: 90972146
#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define W while
#define I inline
#define RI int
#define LL long long
#define Cn const
#define CI Cn int&
#define gc getchar
#define D isdigit(c=gc())
#define pc(c) putchar((c))
using namespace std;
namespace Debug{
	Tp I void _debug(Cn char* f,Ty t){cerr<<f<<'='<<t<<endl;}
	Ts I void _debug(Cn char* f,Ty x,Ar... y){W(*f!=',') cerr<<*f++;cerr<<'='<<x<<",";_debug(f+1,y...);}
	Tp ostream& operator<<(ostream& os,Cn vector<Ty>& V){os<<"[";for(Cn auto& vv:V) os<<vv<<",";os<<"]";return os;}
	#define gdb(...) _debug(#__VA_ARGS__,__VA_ARGS__)
}using namespace Debug;
namespace FastIO{
	Tp I void read(Ty& x){char c;int f=1;x=0;W(!D) f=c^'-'?1:-1;W(x=(x<<3)+(x<<1)+(c&15),D);x*=f;}
	Ts I void read(Ty& x,Ar&... y){read(x),read(y...);}
	Tp I void write(Ty x){x<0&&(pc('-'),x=-x,0),x<10?(pc(x+'0'),0):(write(x/10),pc(x%10+'0'),0);}
	Tp I void writeln(Cn Ty& x){write(x),pc('\n');}
}using namespace FastIO;
Cn int N=2e5+10;
int n,c[N],f[N],sz[N],g[N],p[N],h[N];LL T,Ans[N];
vector<int> G[N];
#define pb push_back
I LL R(CI x){
	return 1LL*x*(x+1)/2;
}
I void DFS(CI x,CI fa){
	RI t=p[c[x]];g[x]=p[c[x]],p[c[x]]=x,sz[x]=1;for(auto i:G[x]) i^fa&&(DFS(i,x),sz[x]+=sz[i],0);
	p[c[x]]=t,f[g[x]]+=sz[x],!g[x]&&(h[c[x]]+=sz[x]);
}
I void DFS2(CI x,CI fa){
	RI t=p[c[x]];for(auto i:G[x]) i^fa&&(p[c[x]]=0,DFS2(i,x),Ans[c[x]]+=R(sz[i]-p[c[x]]),0);
	p[c[x]]=t+sz[x];
}
int main(){
	RI i,x,y;for(read(n),T=1LL*n*(n+1)/2,i=1;i<=n;i++) read(c[i]);
	for(i=1;i<n;i++) read(x,y),G[x].pb(y),G[y].pb(x);
	DFS(1,0),DFS2(1,0);//gdb(f[0]);
	for(i=1;i<=n;i++) /*gdb(i,Ans[i],h[i]),*/i^c[1]&&(Ans[i]+=R(n-h[i]));
	for(i=1;i<=n;i++) writeln(T-Ans[i]);
	return 0;
}

Submission Info

Judge Result