P2664树上游戏(点分治)

 思路:

解决树上路径问题,我们可以想到用点分治.现在问题转化为对树上lca为根节点的点对的路径的颜色进行计数.

直接做比较难,我们考虑贡献.注意到,对于一个点x,如果x的颜色是根节点到x这段路径上第一次出现的颜色,那么x会对根节点除了以x所在分支对应的儿子为子树以外的所有节点(包括根节点)产生sz[x]的贡献,sz[x]指的是以x为根的子树的大小,这里我们不考虑其他节点也出现了这个颜色的情况,我们后面会解决这个问题.接下来具体是每个点怎么计算答案.

1.对于根节点来说,所有跟他不相同颜色都会产生贡献,然后其本身会跟其他点产生贡献,只要把这两个贡献相加即可.

2.对于除了根节点以外的任何一点x,我们由根节点的计算可以得到启发,可以计算别的点对这个点的贡献和这个点对别的点的贡献.首先别的点对这个点的贡献我们上面算出来了,然后我们减去这个点对其他点的贡献,也就是上文提到的sz[x].然后再减去这个点x到根节点中出现过的颜色的贡献,因为在这条路径上已经算过一次了,再算一次就重复了.这个是别的点对这个点的贡献.接下来计算这个点对其他点的贡献,从根节点到当前点x的路径上出现一次的颜色,都会与根节点除了以x所在分支对应的儿子为子树以外的所有节点组成路径,假设这个点所在分支对应根节点的y儿子,从根节点到x上出现的颜色数为num,那么这个贡献就是num * (sz[root] - sz[y]).

更详细的可见代码:
 

#include <bits/stdc++.h>
#define int long long                                         
#define IOS ios::sync_with_stdio(false), cin.tie(0) 
#define ll long long 
#define double long double
#define ull unsigned long long 
#define PII pair<int, int> 
#define PDI pair<double, int> 
#define PDD pair<double, double> 
#define debug(a) cout << #a << " = " << a << endl 
#define point(n) cout << fixed << setprecision(n)
#define all(x) (x).begin(), (x).end() 
#define mem(x, y) memset((x), (y), sizeof(x)) 
#define lbt(x) (x & (-x)) 
#define SZ(x) ((x).size()) 
#define inf 0x3f3f3f3f 
#define INF 0x3f3f3f3f3f3f3f3f
namespace nqio{const unsigned R = 4e5, W = 4e5; char *a, *b, i[R], o[W], *c = o, *d = o + W, h[40], *p = h, y; bool s; struct q{void r(char &x){x = a == b && (b = (a = i) + fread(i, 1, R, stdin), a == b) ? -1 : *a++;} void f(){fwrite(o, 1, c - o, stdout); c = o;} ~q(){f();}void w(char x){*c = x;if (++c == d) f();} q &operator >>(char &x){do r(x);while (x <= 32); return *this;} q &operator >>(char *x){do r(*x); while (*x <= 32); while (*x > 32) r(*++x); *x = 0; return *this;} template<typename t> q&operator>>(t &x){for (r(y),s = 0; !isdigit(y); r(y)) s |= y == 45;if (s) for (x = 0; isdigit(y); r(y)) x = x * 10 - (y ^ 48); else for (x = 0; isdigit(y); r(y)) x = x * 10 + (y ^ 48); return *this;} q &operator <<(char x){w(x);return *this;}q &operator<< (char *x){while (*x) w(*x++); return *this;}q &operator <<(const char *x){while (*x) w(*x++); return *this;}template<typename t> q &operator<< (t x) {if (!x) w(48); else if (x < 0) for (w(45); x; x /= 10) *p++ = 48 | -(x % 10); else for (; x; x /= 10) *p++ = 48 | x % 10; while (p != h) w(*--p);return *this;}}qio; }using nqio::qio;
using namespace std;
const int N = 2e6 + 10;
int n, h[N], v[N], c[N], to[N], tot;
int root, mx[N], sum, sz[N], vis[N], ans[N];
int cnt[N], color[N], now, num, temp;//now表示当前所有节点对颜色的贡献
void add(int a, int b) {
	v[++tot] = b, to[tot] = h[a], h[a] = tot;
}
void find_rt(int x, int fa) {
	mx[x] = 0, sz[x] = 1;
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]] || y == fa) continue;
		find_rt(y, x);
		sz[x] += sz[y];
		mx[x] = max(mx[x], sz[y]);
	}
	mx[x] = max(mx[x], sum - mx[x]);
	if (mx[x] < mx[root]) root = x;
}
void dfs1(int x, int fa) {
	++cnt[c[x]]; sz[x] = 1;
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]] || y == fa) continue;
		dfs1(y, x);
		sz[x] += sz[y];
	}
	if (cnt[c[x]] == 1) {
		now += sz[x];
		color[c[x]] += sz[x];
	}
	--cnt[c[x]];
}
void dfs2(int x, int fa, int val) {
	++cnt[c[x]];
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]] || y == fa) continue;
		dfs2(y, x, val);
	}
	if (cnt[c[x]] == 1) {
		now += sz[x] * val;
		color[c[x]] += sz[x] * val;
	}
	--cnt[c[x]];
}
void dfs3(int x, int fa) {
	++cnt[c[x]];
	if (cnt[c[x]] == 1) {
		now -= color[c[x]];
		++num;		
	}
	ans[x] += now + num * temp;
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]] || y == fa) continue;
		dfs3(y, x);
	}
	if (cnt[c[x]] == 1) {
		now += color[c[x]];
		--num;
	}
	--cnt[c[x]];
}
void clear(int x, int fa) {
	color[c[x]] = cnt[c[x]] = 0;
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]] || y == fa) continue;
		clear(y, x);
	}
}
void calc(int x) {
	dfs1(x, 0);
	ans[x] += now - color[c[x]] + sz[x];
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]]) continue;
		++cnt[c[x]];
		now -= sz[y];
		color[c[x]] -= sz[y];
		dfs2(y, x, -1);
		--cnt[c[x]];
		temp = sz[x] - sz[y];
		dfs3(y, x);
		++cnt[c[x]];
		now += sz[y];
		color[c[x]] += sz[y];
		dfs2(y, x, 1);
		--cnt[c[x]];
	}
	now = num = 0;
	clear(x, 0);
}
void divid(int x) {
	calc(x); vis[x] = 1;
	for (int i = h[x], y; i; i = to[i]) {
		if (vis[y = v[i]]) continue;
		root = 0, sum = sz[y];
		find_rt(y, x);
		divid(root);
	}
}
signed main() {
	qio >> n;
	for (int i = 1; i <= n; ++i) qio >> c[i];
	for (int i = 1; i <  n; ++i) {
		int x, y;
		qio >> x >> y;
		add(x, y), add(y, x);
	}
	mx[root = 0] = n + 1, sum = n; find_rt(1, 0);
	divid(root);
	for (int i = 1; i <= n; ++i) qio << ans[i] << "\n";
}