[洛谷P2664] 树上游戏

前言

本来想练一下点分树来着，结果这题可以 \(O(n)\)，结果还莫名拿了 rk1？？？

我看懂了题解的第一步转化，然后死活看不懂接下来的步骤，被迫只能自己写了/kk

题目

洛谷

讲解

拆贡献是肯定的，但是这里我们不拆点的贡献，而是拆颜色的贡献。

考虑对于一种颜色 col，我们把所有颜色是 col 的点相连的边全部断掉，可以发现每个点的贡献是 n - 自己所在连通块大小(当然颜色是 col 的点除外)。

令总颜色数量为 m，那么每个点的答案就是：\(n\times m-\sum_{col}siz_{col,i}\)，注意点的颜色是 col 时这个东西需要微调。

可以用树上差分简单求解，复杂度 \(O(n)\)。

代码

//12252024832524
#include <bits/stdc++.h>
#define TT template<typename T>
using namespace std;

typedef long long LL;
const int MAXN = 100005;
int n,m;
int col[MAXN];
bool vis[MAXN];
LL ans[MAXN],dp[MAXN],wr[MAXN];

LL Read()
{
	LL x = 0,f = 1; char c = getchar();
	while(c > '9' || c < '0'){if(c == '-') f = -1;c = getchar();}
	while(c >= '0' && c <= '9'){x = (x*10) + (c^48);c = getchar();}
	return x * f;
}
TT void Put1(T x)
{
	if(x > 9) Put1(x/10);
	putchar(x%10^48);
}
TT void Put(T x,char c = -1)
{
	if(x < 0) putchar('-'),x = -x;
	Put1(x); if(c >= 0) putchar(c);
}
TT T Max(T x,T y){return x > y ? x : y;}
TT T Min(T x,T y){return x < y ? x : y;}
TT T Abs(T x){return x < 0 ? -x : x;}

int head[MAXN],tot,hd[MAXN];
struct edge{
	int v,nxt;
}e[MAXN<<2];
void Add_Edge(int u,int v){
	e[++tot] = edge{v,head[u]};
	head[u] = tot;
}
void Add_Double_Edge(int u,int v){
	Add_Edge(u,v);
	Add_Edge(v,u);
}

int f[MAXN],siz[MAXN],fff[MAXN];
void dfs1(int x,int fa) {
	siz[x] = 1;
	e[++tot] = edge{x,hd[f[col[x]]]}; hd[f[col[x]]] = tot;
	int tmp = f[col[x]]; f[col[x]] = x;
	for(int i = head[x],v; i ;i = e[i].nxt){
		if((v = e[i].v) == fa) continue;
		LL lst = wr[x],val;
		dfs1(v,x);
		val = (!x ? 1ll * n * m : siz[v]) - (wr[x]-lst);
		dp[v] += val; //printf("add %d %lld\n",v,val);
		for(int j = hd[x]; j ;j = hd[x] = e[j].nxt){
			if(!x) dp[e[j].v] -= n-fff[col[e[j].v]];
			else dp[e[j].v] -= val;
		}
		siz[x] += siz[v];
	}
	f[col[x]] = tmp;
	if(!f[col[x]]) fff[col[x]] += siz[x];
	wr[f[col[x]]] += siz[x];
}
void dfs2(int x,int fa){
	ans[x] = dp[x] += dp[fa];
	for(int i = head[x],v; i ;i = e[i].nxt){
		if((v = e[i].v) == fa) continue;
		dfs2(v,x);
	}
}

int main()
{
//	freopen(".in","r",stdin);
//	freopen(".out","w",stdout);
	n = Read();
	for(int i = 1;i <= n;++ i){
		col[i] = Read();
		if(!vis[col[i]]) vis[col[i]] = 1,++m;
	}
	Add_Edge(0,1);
	for(int i = 1;i < n;++ i) Add_Double_Edge(Read(),Read());
	dfs1(0,0);
	dfs2(1,0);
	for(int i = 1;i <= n;++ i) Put(1ll*n*m-ans[i]-n,'\n');
	return 0;
}