【Luogu P2664】树上游戏（点分治）

题意

一棵树，每个点有颜色。s(i,j)表示i到j路径上的颜色个数。对每一个点$i$求${\sum }_{j=1}^{n}s(i,j)$

思路

点分治简单题。不过作为恢复训练还是不错的。

考虑在重心为$zx$的树中被贡献的点为$u$的情况，这棵点分出来的树上每个点$v$将如何贡献。

1. $v$和$u$在同一个子树内：无法贡献。
2. $v$和$u$不在同一棵子树，$u$到$v$的路径上也没有和$v$同色的点：贡献为$v$的子树大小。
3. $v$和$u$不在同一棵子树，$u$到$v$的路径上有和$v$同色的点：无法贡献。

简单理解2和3就是假如路上已经有一个同色点了，那么经过这个点的所有路径必然经过那个路上的同色点。

所以我们再把路径拆成$u$到$zx$和$zx$到$v$。先处理处每个$v$对$zx$的贡献$con[v]$。然后dfs每棵子树，$u$到$zx$的路径上所有颜色都会出现在所有$n-siz[rt[u]]$（n是点分出的树大小，siz[rt[u]]是u这棵子树的大小）中，在dfs和回溯的时候维护贡献就好了。

复杂度$O(nlogn)$，常数有点大。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 1e5 + 10;
const int M = N << 1;
const int inf = 1e9 + 7;
namespace Graph
{
	int h[N], ecnt, nxt[M], v[M];
	void clear(){ecnt = 1;}
	void add_dir(int _u, int _v){
	    v[++ecnt] = _v;
	    nxt[ecnt] = h[_u]; h[_u] = ecnt;
	}
	void add_undir(int _u, int _v){
	    add_dir(_u, _v);
	    add_dir(_v, _u);
	}
}
using namespace Graph;
int n, c[N], val[N], tmp[N];
LL con[N], totcon, ans[N];
int zx, mn, sz, siz[N];
bool vis[N];

void get_zx(int u, int fa)
{
	int son = 0;
	siz[u] = 1;
	for (int i = h[u]; i; i = nxt[i])
		if (v[i] != fa && !vis[v[i]]){
			get_zx(v[i], u);
			siz[u] += siz[v[i]];
			son = max(son, siz[v[i]]);
		}
	son = max(son, sz - siz[u]);
	if (son < mn) mn = son, zx = u;
}

void get_val(int u, int fa)
{
	bool hav = tmp[c[u]] == 0;
	++ tmp[c[u]];
	siz[u] = 1;
	for (int i = h[u]; i; i = nxt[i])
		if (v[i] != fa && !vis[v[i]]){
			get_val(v[i], u);
			siz[u] += siz[v[i]];
		}
	-- tmp[c[u]];
	val[u] = hav ? siz[u] : 0;
	con[c[u]] += val[u];
	totcon += val[u];
}

void del_con(int u, int fa, bool fl)
{
	if (fl){
		totcon -= val[u];
		con[c[u]] -= val[u];
	}
	else{
		totcon += val[u];
		con[c[u]] += val[u];
	}
	for (int i = h[u]; i; i = nxt[i])
		if (v[i] != fa && !vis[v[i]])
			del_con(v[i], u, fl);
}

void calc(int u, int fa, int _sz)
{
	if (val[u] > 0) totcon += _sz - con[c[u]];
	ans[u] += totcon;
	for (int i = h[u]; i; i = nxt[i])
		if (v[i] != fa && !vis[v[i]])
			calc(v[i], u, _sz);
	if (val[u] > 0) totcon -= _sz - con[c[u]];
}

void solve(int _zx, int _sz)
{
	vis[_zx] = 1;
	++tmp[c[_zx]];
	for (int i = h[_zx]; i; i = nxt[i])
		if (!vis[v[i]])
			get_val(v[i], _zx);
	--tmp[c[_zx]];
	ans[_zx] += _sz + totcon;
	for (int i = h[_zx]; i; i = nxt[i])
		if (!vis[v[i]]){
			del_con(v[i], _zx, 1);
			totcon += _sz - siz[v[i]];
			calc(v[i], _zx, _sz - siz[v[i]]);
			totcon -= _sz - siz[v[i]];
			del_con(v[i], _zx, 0);
		}
	for (int i = h[_zx]; i; i = nxt[i])
		if (!vis[v[i]])
			del_con(v[i], _zx, 1);
	for (int i = h[_zx]; i; i = nxt[i])
		if (!vis[v[i]]){
			sz = siz[v[i]]; mn = inf;
			get_zx(v[i], 0);
			solve(zx, sz);
		}
}

int main()
{
	scanf("%d", &n);
	for (int i = 1; i <= n; ++ i)
		scanf("%d", &c[i]);
	clear();
	for (int i = 1; i < n; ++ i){
		int x, y;
		scanf("%d%d", &x, &y);
		add_undir(x, y);
	}
	sz = n; mn = inf;
	get_zx(1, 0);
	solve(zx, n);
	for (int i = 1; i <= n; ++ i)
		printf("%lld\n", ans[i]);
	return 0;
}