HDU 5977：Garden of Eden（点分治 + 状压预处理）

题意：在一棵有n个点的树上，每个点有一个点权，最多有 k 种点权，问 简单路径上点权种类总共为 k 的点对有多少对。


题解：考虑点分治，由于k比较小，用一个state数组记录当前点到分治中心的路径上颜色的种类集合，这个种类集合可以用二进制状压。

每dfs完一棵子树，枚举已经dfs过的子树上的点暴力合并，若两个点的状态合并之后集合大小 = k，则更新答案。这样做是 $n^{2}log(n^2)$

开一个数组 $num$ 用来记录前i棵子树 到 分治中心 路径上种类状态为 i （二进制状压）的点的个数。对每颗子树，枚举 $2^k$个状态，若能合并成$(1<<k-1)$，则更新答案，这样做的复杂度是 $2^k\ast n\ast logn$,已经足够通过此题，但还可以优化。

对于每一个状态，能与之相或得到$(1<<k)-1$的状态数并不是太多，预处理每个状态能合并得到答案的状态，减少每次合并答案的枚举量。

#include
using namespace std;
const int maxm = 1e6 + 10;
const int maxn = 1e5 + 10;
int head[maxm],to[maxm],nxt[maxm],cnt;
int root,sz[maxn],f[maxn],tot,a[maxn];
long long res;
bool done[maxn];
int n,k,state[maxn],num[maxn];
vector<int> t,tmp;
vector<int> g[1500];
void init() {
	res = cnt = 0;
	tmp.clear();t.clear();
	fill(head,head + n + 1,-1);
	fill(done,done + n + 1,0);
	for(int i = 0; i < (1 << k); i++) g[i].clear();
}
void add(int u,int v) {
	to[cnt] = v;
	nxt[cnt] = head[u];
	head[u] = cnt++;
}
void getroot(int u,int fa) {
	sz[u] = 1;f[u] = 0;
	for(int i = head[u]; i + 1; i = nxt[i]) {
		if(to[i] == fa || done[to[i]]) continue;
		getroot(to[i],u);
		sz[u] += sz[to[i]];
		f[u] = max(f[u],sz[to[i]]);
	}
	f[u] = max(f[u],tot - sz[u]);
	if(!root || f[u] < f[root]) root = u;
}
void dfs(int u,int fa) {
	t.push_back(u);tmp.push_back(u);
	for(int i = head[u]; i + 1; i = nxt[i]) {
		if(to[i] == fa || done[to[i]]) continue;
		state[to[i]] = (state[u] | (1 << a[to[i]]));
		dfs(to[i],u);
	}
}
long long solve(int u) {
	long long ans = 0;
	done[u] = true;tmp.clear();
	tmp.push_back(u);state[u] = (1 << a[u]);
	if(state[u] == (1 << k) - 1) ans++;
	num[state[u]] = 1;
	for(int i = head[u]; i + 1; i = nxt[i]) {
		if(done[to[i]]) continue;
		t.clear();
		state[to[i]] = ((1 << a[to[i]]) | state[u]);
		dfs(to[i],u);
		for(auto i : t)
			for(auto j: g[state[i]])
				ans += num[j] * 2;
		for(auto i : t)
			num[state[i]]++;
	}
	for(auto i: tmp)
		num[state[i]] = 0;
	return ans;
}
void divide(int rt) {
	res += solve(rt);
	for(int i = head[rt]; i + 1; i = nxt[i]) {
		if(done[to[i]]) continue;
		root = 0;tot = sz[to[i]];
		getroot(to[i],rt);
		divide(root);
	}
}
int main() {
	while(~scanf("%d%d",&n,&k)) {
		init();
		for(int i = 0; i < (1 << k); i++)
			for(int j = 0; j < (1 << k); j++)
				if((i | j) == ((1 << k) - 1)) g[i].push_back(j);
		for(int i = 1; i <= n; i++) {
			scanf("%d",&a[i]);
			a[i]--;
		}
		for(int i = 1; i < n; i++) {
			int u,v;
			scanf("%d%d",&u,&v);
			add(u,v);add(v,u);
		}
		tot = n;root = 0;
		getroot(1,-1);
		divide(root);
		printf("%lld\n",res);
	}
	return 0;
}