点分治

题目点这里

大连2016 的 G 题，调了半天的点分。

题目为给你一棵树，树上的每个点都有一个权值，求包含所有权值的路径数目。

由于权值最大只有 $10$，所以点分治的思路很明显了，枚举每条路径起点的同时枚举权值找终点，时间复杂度 $O(2^{k}n\log n)$。

#include <bits/stdc++.h>

const int MAXN = 5e4 + 5;

struct Edge {
	struct Node *to;
	struct Edge *next;

	Edge(Node *to, Edge *next) : to(to), next(next) {}
	Edge() {}
} _pool[MAXN * 2], *_cur;

struct Node {
	Edge *fe;
	int sz, box, msz, num;
	bool vis;

	void push(Node *to) {
		fe = new (_cur++)Edge(to, fe);
	}
} node[MAXN];

int n, k, tt;
long long ans;
int cnt[1024];

void getroot(Node *x, Node *fa, Node *&root, int tot) {
	x->sz = 1, x->msz = 0;
	for(Edge *e = x->fe; e; e = e->next) {
		if(e->to->vis || e->to == fa) continue;
		getroot(e->to, x, root, tot);
		x->sz += e->to->sz;
		x->msz = std::max(x->msz, e->to->sz);
	}
	x->msz = std::max(x->msz, tot - x->sz);
	if(!root || x->msz < root->msz) root = x;
}

long long dfs(Node *x, Node *fa) {
	long long ans = 0;
	for(int i = 0; i <= tt; i++) {
		if((i | x->box) == tt) ans += cnt[i];
	}
	cnt[x->box]++;
	for(Edge *e = x->fe; e; e = e->next) {
		if(e->to->vis || e->to == fa) continue;
		e->to->box = (x->box | (1 << (e->to->num - 1)));
		ans += dfs(e->to, x);
	}
	return ans;
}

long long calc(Node *x, Node *fa, int arg) {
	memset(cnt, 0, sizeof(cnt));
	x->box = (arg | (1 << (x->num - 1)));
	return dfs(x, 0);
}

void divid(Node *x) {
	ans = ans + 2 * calc(x, 0, 0);
	x->vis = 1;
	for(Edge *e = x->fe; e; e = e->next) {
		if(e->to->vis) continue;
		ans = ans - 2 * calc(e->to, 0, e->to->box);
		int tot = e->to->sz;
		if(tot < k) continue;
		Node *root = 0;
		getroot(e->to, 0, root, tot);
		divid(root);
	}
	x->vis = 0;
}

void solve() {
	tt = (1 << k) - 1;
	_cur = _pool;
	ans = 0;
	memset(node, 0, sizeof(node));
	for(int i = 1; i <= n; i++) scanf("%d", &node[i].num);
	for(int i = 1; i < n; i++) {
		int u, v;
		scanf("%d%d", &u, &v);
		node[u].push(node + v);
		node[v].push(node + u);
	}
	Node *root = 0;
	getroot(node + 1, 0, root, n);
	divid(root);
	if(k == 1) ans += n;
	printf("%lld\n", ans);
}

int main() {
	while(~scanf("%d%d", &n, &k)) solve();
	return 0;
}