大连2016 的 G 题,调了半天的点分。
题目为给你一棵树,树上的每个点都有一个权值,求包含所有权值的路径数目。
由于权值最大只有 10 10 10,所以点分治的思路很明显了,枚举每条路径起点的同时枚举权值找终点,时间复杂度 O ( 2 k n log n ) O(2^kn\log n) O(2knlogn)。
#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;
}