求调
查看原帖
求调
379420
Xuejiama1227楼主2023/7/24 13:48
// Problem: P2664 树上游戏
// Contest: Luogu
// URL: https://www.luogu.com.cn/problem/P2664
// Memory Limit: 125 MB
// Time Limit: 1000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include<bits/stdc++.h>
using namespace std;
template<typename T>
void read(T&x){
	x=0;char c=getchar();T f=1;
	for(;c<'0'||c>'9';c=getchar())if(c=='-')f=-1;
	for(;c>='0'&&c<='9';c=getchar())x=(x<<3)+(x<<1)+(c&15);
	x=x*f;
}
template<typename T,typename...Args>
void read(T&x,Args&...args){read(x);read(args...);}
const int N=2e5+5;
const int INF=2e9;
typedef long long ll;
int n,nn,rt;
struct node{
	bool vis;
	int cl,siz,mx;
	vector<int>e;
}a[N];
void calc_siz(int x,int fa){
	a[x].siz=1;
	for(auto y:a[x].e)if(!a[y].vis&&y!=fa){
		calc_siz(y,x);
		a[x].mx=max(a[x].mx,nn-a[y].siz);
		a[x].siz+=a[y].siz;
	}
	a[x].mx=max(a[x].mx,nn-a[x].siz);
	if(a[x].mx<a[rt].mx)rt=x;
}
ll ans[N],sum;
int cnt[N],v[N],nrt;
void calc_dis(int x,int fa,int nw){
	a[x].siz=1;
	if(!v[a[x].cl])sum-=cnt[a[x].cl],nw++;
	v[a[x].cl]++;
	ans[x]+=sum+nw*a[nrt].siz;
	for(auto y:a[x].e)if(!a[y].vis&&y!=fa){
		calc_dis(y,x,nw);
		a[x].siz+=a[y].siz;
    }
	v[a[x].cl]--;
	if(!v[a[x].cl])sum+=cnt[a[x].cl];
}
void calc_cnt(int x,int fa){
	if(!v[a[x].cl]){
		cnt[a[x].cl]+=a[x].siz;
		sum+=a[x].siz;
	}
	v[a[x].cl]++;
	for(auto y:a[x].e)if(!a[y].vis&&y!=fa)calc_cnt(y,x);
	v[a[x].cl]--;
}
void clear(int x,int fa,int nw){
	if(!v[a[x].cl])nw++;
	v[a[x].cl]++;ans[x]-=nw;ans[nrt]+=nw;
	for(auto y:a[x].e)if(!a[y].vis&&y!=fa)clear(y,x,nw);
	v[a[x].cl]--;cnt[a[x].cl]=0;
}
void clear2(int x,int fa){
	cnt[a[x].cl]=0;
	for(auto y:a[x].e)if(!a[y].vis&&y!=fa)clear2(y,x);
}
int sn[N];
void divid(int x){
	a[x].vis=1;nrt=x;ans[x]++;
	int tot=0,i,y;
	for(auto y:a[x].e)if(!a[y].vis)sn[++tot]=y;
	a[x].siz=sum=cnt[a[x].cl]=1;
	v[a[x].cl]++;
	for(i=1;i<=tot;i++){
    	y=sn[i];
    	calc_dis(y,x,0);
		calc_cnt(y,x);
		a[x].siz+=a[y].siz;
		cnt[a[x].cl]+=a[y].siz;
		sum+=a[y].siz;
	}
	clear2(x,0);
	a[x].siz=sum=cnt[a[x].cl]=1;
	for(i=tot;i>=1;i--){
    	y=sn[i];
    	calc_dis(y,x,0);
    	calc_cnt(y,x);
		a[x].siz+=a[y].siz;
		cnt[a[x].cl]+=a[y].siz;
		sum+=a[y].siz;
	}
	v[a[x].cl]--;
	clear(x,0,0);
	for(auto y:a[x].e)if(!a[y].vis){
		nn =a[y].siz,rt=0,a[rt].mx=n+1;
		calc_siz(y,x);
		divid(rt);
	}
}
signed main(){
	int i,x,y;
	read(n);
	for(i=1;i<=n;i++)read(a[i].cl);
	for(i=1;i<n;i++){
		read(x,y);
		a[x].e.push_back(y);
		a[y].e.push_back(x);
	}
	rt=0,nn=n,a[rt].mx=n+1;
	calc_siz(1,0);
	divid(rt);
	for(i=1;i<=n;i++)printf("%lld\n",ans[i]);
	return 0;
}

T四个点

2023/7/24 13:48
加载中...