[LuoGu]P2664 树上游戏

Portal

这题真的好.

看到树上路径, 脑子里就要点分治

这一题对于每个点都要计算一遍, 如果暴算实在不好算, 这样我们就可以考虑算贡献.

直接计算每种颜色的贡献.

因为一条过重心的路径中, 可能两边都会有相同颜色, 那么我们就只考虑当前点到分治中心的链上.

然后把多算的这条链上的东西删除就可以了.

Copy From Here

那么我们就可以这样做了：

１.对树进行第一遍dfs，预处理size和上方性质中的贡献。（开一个color数组即可记录贡献），同时记录贡献总和sum

２.枚举跟的所有子树，先把子树扫一遍清除其在color数组中的所有贡献。接着，对于该子树中的每一个点ｊ：

设Ｘ＝sigma color[j 到根上（不包括根）的所有颜色] 由于这些颜色已经出现过，我们不能在该子树外计算其贡献）

设num为ｊ到根上（不包括根）的颜色数

设Ｙ为size[root]-size[该子树]（即所有其他子树＋根的点数）

则ans[j]+＝sum-X+num*Y;

３.别忘了计算单独ｒｏｏｔ的ａｎｓ

　　ans[root]+=sum-color[根的颜色]+size[root]

４．清空贡献数组以及其他东西

Code

#include<bits/stdc++.h>
using namespace std;
#define rep(i, a, b) for(int i = (a), i##_end_ = (b); i <= i##_end_; ++i)
#define drep(i, a, b) for(int i = (a), i##_end_ = (b); i >= i##_end_; --i)
#define clar(a, b) memset((a), (b), sizeof(a))
#define debug(...) fprintf(stderr, __VA_ARGS__)
typedef long long LL;
typedef long double LD;
int read() {
    char ch = getchar();
    int x = 0, flag = 1;
    for (;!isdigit(ch); ch = getchar()) if (ch == '-') flag *= -1;
    for (;isdigit(ch); ch = getchar()) x = x * 10 + ch - 48;
    return x * flag;
}
void write(int x) {
    if (x < 0) putchar('-'), x = -x;
    if (x >= 10) write(x / 10);
    putchar(x % 10 + 48);
}

const int Maxn = 100009;

struct edge {
	int to, nxt;
}g[Maxn << 1];

int n, a[Maxn], head[Maxn], e;

void add(int u, int v) {
	g[++e] = (edge){v, head[u]}, head[u] = e;
}

void init() {
	clar(head, -1);
	n = read();
	rep (i, 1, n) a[i] = read();
	rep (i, 1, n - 1) {
		int u = read(), v = read();
		add(u, v), add(v, u);
	}
}

LL ans[Maxn], vis[Maxn], Tmp, heart;
int size[Maxn];

void dfs(int u, int pa) { 
	size[u] = 1;
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (v != pa && !vis[v]) {
			dfs(v, u);
			size[u] += size[v];
		}
	}
}

void centroid(int u, int pa, int field) {
	int bbs = field - size[u];
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (!vis[v] && v != pa) {
			centroid(v, u, field);
			bbs = max(bbs, size[v]);
		}
	}
	if (bbs < Tmp) Tmp = bbs, heart = u;
}

LL color[Maxn], Sum, cnt[Maxn], size1[Maxn];

void dfs_init(int u, int pa) {
	size1[u] = 1; ++cnt[a[u]];
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (v != pa && !vis[v]) {
			dfs_init(v, u);
			size1[u] += size1[v];
		}
	}

	if (cnt[a[u]] == 1) {
		color[a[u]] += size1[u];
		Sum += size1[u];
	}
	--cnt[a[u]];
}

void change(int u, int pa, int delta) {
	++cnt[a[u]];
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (v != pa && !vis[v]) change(v, u, delta);
	}

	if (cnt[a[u]] == 1) {
		color[a[u]] += size1[u] * 1ll * delta;
		Sum += size1[u] * 1ll * delta;
	}
	--cnt[a[u]];
}

void count(int u, int pa, int X, int Y, int num) {
	++cnt[a[u]];
	if (cnt[a[u]] == 1) {
		X += color[a[u]];
		++Y;
	}
	ans[u] += Sum - X + num * Y;
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (!vis[v] && v != pa) count(v, u, X, Y, num);
	}
	--cnt[a[u]];
}

void dfs_clear(int u, int pa) {
	color[a[u]] = 0, cnt[a[u]] = 0, size1[u] = 0;
	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (v != pa && !vis[v]) dfs_clear(v, u);
	}
}

void calc(int u) {
	dfs_init(u, 0);
	ans[u] += Sum - color[a[u]] + size1[u];

	for (int i = head[u]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (!vis[v]) {
			++cnt[a[u]];
			Sum -= size1[v]; color[a[u]] -= size1[v];	
			change(v, u, -1);
			--cnt[a[u]];
			count(v, u, 0, 0, size1[u] - size1[v]);
			++cnt[a[u]];
			Sum += size1[v]; color[a[u]] += size1[v];	
			change(v, u, 1);
			--cnt[a[u]];
		}
	}

	dfs_clear(u, 0), Sum = 0;
}

void divConquer(int u) {
	dfs(u, 0), Tmp = INT_MAX;
	centroid(u, 0, size[u]);
	calc(heart); vis[heart] = 1;
	for (int i = head[heart]; ~i; i = g[i].nxt) {
		int v = g[i].to;
		if (!vis[v]) divConquer(v);
	}
}

void solve() {
	divConquer(1);
	rep (i, 1, n) printf("%lld\n", ans[i]);
}

int main() {
	freopen("LG2664.in", "r", stdin);
	freopen("LG2664.out", "w", stdout);

	init();
	solve();

#ifdef Qrsikno
    debug("\nRunning time: %.3lf(s)\n", clock() * 1.0 / CLOCKS_PER_SEC);
#endif
    return 0;
}