Luogu P2664 树上游戏

Link

妙妙题。

看见树上路径求和，直接点分治。

对于某个颜色对根节点的贡献，显然是所有距离根节点最近的点的子树和。

考虑其他子树对当前子树的贡献，分成两部分：

1. 到子树外才出现的颜色

2. 子树内出现的颜色

维护 \(cols_i\) 为颜色 \(i\) 最早出现的所有结点的 \(siz\)，给子树算贡献的时候先把子树部分对 \(cols_i\) 的贡献删掉，然后再算（记得加回来）。

复杂度 \(O(n\log n)\)

#include <stdio.h>
#include <algorithm>
#include <string.h>
#include <cctype>
using namespace std;
typedef long long ll;
const int INF=0x3f3f3f3f;
const int mo=1e9+7;
inline int read(){
	char ch=getchar();int x=0, f=1;
	while(!isdigit(ch)){if(ch=='-') f=-1; ch=getchar();}
	while(isdigit(ch)){x=(x<<3)+(x<<1)+ch-'0';ch=getchar();}
	return x*f;
}
inline void write(int x){
    if(x<0) putchar('-'), x=-x;
    if(x>9) write(x/10);
    putchar(x%10+'0');
}
inline int ksm(int a, int b){
	int ret=1;
	for(; b; b>>=1, a=1ll*a*a%mo)
		if(b&1) ret=1ll*ret*a%mo;
	return ret;
}
const int N=1e5+5;
int tot, h[N];
struct Edge{int to, nxt;}d[N*2];
void add(int x, int y){d[++tot]=(Edge){y, h[x]};h[x]=tot;}
int n, col[N], colc[N], S, cols[N], siz[N];
int G, Mn, ans[N], idc, vis[N];
void GetGra(int x ,int up){
	vis[x]=idc, siz[x]=1;int mx=0;
	for(int i=h[x], v; i; i=d[i].nxt)
		if(vis[v=d[i].to]!=-1&&vis[v]<idc)
			GetGra(v, up), siz[x]+=siz[v], 
			mx=max(mx, siz[v]);
	mx=max(mx, up-siz[x]);
	if(mx<=Mn) Mn=mx, G=x;
}
int bin[N], top;
void PreTre(int x, int flg){
	vis[x]=idc, colc[col[x]]++;
	if(flg==0) siz[x]=1;//处理根
	for(int i=h[x], v; i; i=d[i].nxt)
		if(vis[v=d[i].to]<idc&&vis[v]!=-1)
			PreTre(v, flg), siz[x]+=(flg?0:siz[v]);
	if(colc[col[x]]==1){
		int typ=flg?flg:1;
		S+=typ*siz[x], cols[col[x]]+=typ*siz[x];
		if(!flg) bin[++top]=col[x];
	}
	colc[col[x]]--;
}
int Els;
void Calc(int x, int Gx){
	colc[col[x]]++, vis[x]=idc;
    if(colc[col[x]]==1)
        Gx=Gx-cols[col[x]]+Els;
        //如果遇到了第一个同颜色的
        //说明不需要子树外面的点满足颜色就珂以
        //换句话来说就是 x 子树内只要过了 x 就有它颜色的贡献
        //此时子树外的点全都不用考虑颜色都有贡献
    ans[x]+=Gx;
	for(int i=h[x], v; i; i=d[i].nxt)
		if(vis[v=d[i].to]<idc&&vis[v]!=-1)
			Calc(v, Gx);
	colc[col[x]]--;
}
void dfz(int rt){
	top=0, idc++, S=0;PreTre(rt, 0), vis[rt]=-1;
	ans[rt]+=S, S-=cols[col[rt]], colc[col[rt]]=1;
	for(int i=h[rt], v; i; i=d[i].nxt)
		if(vis[v=d[i].to]!=-1){
			idc++, PreTre(v, -1);//消去贡献
			Els=siz[rt]-siz[v];
			idc++, Calc(v, S+Els);
			//因为开始去掉乐根节点颜色贡献所以要加上
			//根节点的贡献 Els
			idc++, PreTre(v, 1);//加回来
		}
	colc[col[rt]]=0;
	for(int i=1; i<=top; i++) cols[bin[i]]=0;
	for(int i=h[rt], v; i; i=d[i].nxt)
		if(vis[v=d[i].to]!=-1)
			idc++, G=v, Mn=siz[v],
			GetGra(v, Mn), dfz(G);
}
signed main(){
	n=read();
	for(int i=1; i<=n; i++) col[i]=read();
	for(int i=1, u, v; i<n; i++)
		u=read(), v=read(),
		add(u, v), add(v, u);
	idc=1, G=1, Mn=n,
	GetGra(1, n), dfz(G);
	for(int i=1; i<=n; i++) printf("%d\n", ans[i]);
	return 0;
}