题意:在一棵有n个点的树上,每个点有一个点权,最多有 k 种点权,问 简单路径上点权种类总共为 k 的点对有多少对。
题解:考虑点分治,由于k比较小,用一个state数组记录当前点到分治中心的路径上颜色的种类集合,这个种类集合可以用二进制状压。
每dfs完一棵子树,枚举已经dfs过的子树上的点暴力合并,若两个点的状态合并之后集合大小 = k,则更新答案。这样做是 n 2 l o g ( n 2 ) n^2log(n^2) n2log(n2)
开一个数组 n u m num num 用来记录前i棵子树 到 分治中心 路径上种类状态为 i (二进制状压)的点的个数。对每颗子树,枚举 2 k 2^k 2k个状态,若能合并成 ( 1 < < k − 1 ) (1 << k - 1) (1<<k−1),则更新答案,这样做的复杂度是 2 k ∗ n ∗ l o g n 2^k * n * logn 2k∗n∗logn,已经足够通过此题,但还可以优化。
对于每一个状态,能与之相或得到 ( 1 < < k ) − 1 (1<< k) - 1 (1<<k)−1的状态数并不是太多,预处理每个状态能合并得到答案的状态,减少每次合并答案的枚举量。
#include<bits/stdc++.h>
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;
}