洛谷P2664 树上游戏 【点分治 + 差分】

题目

lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义s(i,j) 为i 到j 的颜色数量。以及

现在他想让你求出所有的sum[i]

输入格式

第一行为一个整数n，表示树节点的数量

第二行为n个整数，分别表示n个节点的颜色c[1],c[2]……c[n]

接下来n-1行，每行为两个整数x,y，表示x和y之间有一条边

输出格式

输出n行，第i行为sum[i]

输入样例

5

1 2 3 2 3

1 2

2 3

2 4

1 5

输出样例

10

9

11

9

12

提示

sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10

sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9

sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11

sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9

sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12

对于40%的数据，n<=2000

对于100%的数据，1<=n,c[i]<=10^5

题解

明显点分治即可

对于每棵分治出来的树，考虑过根的所有路径对树内点的影响

首先单独考虑一种颜色的影响，从根节点出发到每棵子树的每个点\(u\)，\(u\)节点在该颜色下会产生贡献当且仅当\(u\)到根的路径上有该颜色的节点

所以我们只要找出一个子树中所有颜色为该颜色，且其祖先中没有该颜色【也就是最高的该颜色点】，其子树所有点都会产生贡献，那么所有的对根的贡献就是所有这样点的子树大小之和

考虑对子树内的点，就减去该子树的贡献，就转化为和根类似的了

每当第一次经过一种颜色的点时，其子树内所有点经过该点必定产生该颜色的贡献，此时把该颜色的贡献改为剩余子树的大小即可

还有，根节点的颜色特殊考虑

不知讲清楚没有，仔细想想还是很明显的

不过写起来细节真的多

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LL long long int
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define REP(i,n) for (int i = 1; i <= (n); i++)
using namespace std;
const int maxn = 100005,maxm = 200005,INF = 1000000000;
inline int read(){
	int out = 0,flag = 1; char c = getchar();
	while (c < 48 || c > 57){if (c == '-') flag = -1; c = getchar();}
	while (c >= 48 && c <= 57){out = (out << 3) + (out << 1) + c - 48; c = getchar();}
	return out * flag;
}
int h[maxn],ne = 2;
struct EDGE{int to,nxt;}ed[maxm];
inline void build(int u,int v){
	ed[ne] = (EDGE){v,h[u]}; h[u] = ne++;
	ed[ne] = (EDGE){u,h[v]}; h[v] = ne++;
}
int n,c[maxn],F[maxn],N,Siz[maxn],rt,vis[maxn],fa[maxn];
int id[maxn],st[maxn],top,Vis[maxn],now,tot,tots;
LL D[maxn],Sum[maxn],sum[maxn],Sumt[maxn],ttt;
void getrt(int u){
	Siz[u] = 1; F[u] = 0;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]){
		fa[to] = u; getrt(to);
		Siz[u] += Siz[to];
		F[u] = max(F[u],Siz[to]);
	}
	F[u] = max(F[u],N - Siz[u]);
	if (F[u] < F[rt]) rt = u;
}
void dfs1(int u){
	Siz[u] = 1;
	if (!id[c[u]]) st[++top] = c[u],id[c[u]] = top;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]){
		fa[to] = u; dfs1(to);
		Siz[u] += Siz[to];
	}
}
void dfs2(int u){
	int p = id[c[u]],flag = 0;
	if (p != 1 && Vis[p] != now) Sum[p] += Siz[u],Vis[p] = now,flag = 1;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs2(to);
	if (flag) Vis[p] = 0;
}
void dfs3(int u){
	int p = id[c[u]],flag = 0;
	if (p != 1 && Vis[p] != now) Sumt[p] += Siz[u],ttt += Siz[u],Vis[p] = now,flag = 1;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs3(to);
	if (flag) Vis[p] = 0;
}
void dfs4(int u){
	int p = id[c[u]],flag = 0;
	D[u] = D[fa[u]];
	if (p != 1 && Vis[p] != now){
		D[u] -= (Sum[p] - Sumt[p]) - (tot - tots);
		Vis[p] = now; flag = 1;
	}
	sum[u] += D[u] + (tot - tots);
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs4(to);
	if (flag) Vis[p] = 0;
}
void dfs5(int u){
	Sumt[id[c[u]]] = 0;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs5(to);
}
void dfs6(int u){
	D[u] = 0;
	Redge(u) if (!vis[to = ed[k].to] && to != fa[u]) dfs6(to);
}
void solve(int u){
	vis[u] = true; Siz[u] = 1;
	st[top = 1] = c[u]; id[c[u]] = top;
	Redge(u) if (!vis[to = ed[k].to]){
		fa[to] = u; dfs1(to);
		Siz[u] += Siz[to];
	}
	now = 0; tot = Siz[u];
	Redge(u) if (!vis[to = ed[k].to]){
		now++; dfs2(to);
	}
	REP(i,top) D[u] += Sum[i];
	sum[u] += D[u] + Siz[u];
	Redge(u) if (!vis[to = ed[k].to]){
		now++; ttt = 0; tots = Siz[to]; dfs3(to);
		now++; D[u] -= ttt; dfs4(to);
		D[u] += ttt; now++; dfs5(to);
	}
	D[u] = 0;
	Redge(u) if (!vis[to = ed[k].to]) dfs6(to);
	REP(i,top) Vis[i] = Sum[i] = Sumt[i] = id[st[i]] = 0;
	Redge(u) if (!vis[to = ed[k].to]){
		N = Siz[to]; F[rt = 0] = INF;
		getrt(to); solve(rt);
	}
}
int main(){
	n = read();
	REP(i,n) c[i] = read();
	for (int i = 1; i < n; i++) build(read(),read());
	F[rt = 0] = INF; N = n;
	getrt(1); solve(rt);
	REP(i,n) printf("%lld\n",sum[i]);
	return 0;
}