题目
题意
求树上满足条件的点对的个数,条件是在这两点间不经过重复的点能吃到的苹果类型的总数大于 k。
题解
这道题和Tree 这道题相似,用树的分治的方法来做。但 Tree 算的是距离,这道题算的是苹果种类数,所以可以用二进制的方法将状态压缩一下。
还有一个地方需要注意的是,在进行路径合并的时候,需要用子集枚举的方法找到满足要求的所有路径。
代码
#include <algorithm>
#include <bitset>
#include <cstring>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <climits>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
const int MAX = 50005;
vector<int>G[MAX];
int vis[MAX];
int colour[MAX];
int pre[11];
long long req;
int getSize(int u,int fa){
int len = G[u].size();
int sum = 1;
for(int i=0;i<len;++i){
int v = G[u][i];
if(v != fa && vis[v] == 0){
sum += getSize(v,u);
}
}
return sum;
}
int getRoot(int u, int fa,int &root, int &mx, int n) {
int len = G[u].size();
int sum = 1;
int mx1 = 0;
for(int i=0;i<len;++i){
int v = G[u][i];
if(v != fa && vis[v] == 0){
int temp = getRoot(v,u,root,mx,n);
mx1= max(mx1, temp);
sum += temp;
}
}
mx1 = max(mx1, (n - sum));
if(mx1 < mx){
root = u;
mx = mx1;
}
return sum;
}
long long sum;
int N[MAX];
int num;
int n,k;
long long Hash[5000];
void getDeep(int u, int fa, int d){
d |= pre[colour[u]];
N[num++] = d;
int len = G[u].size();
for(int i=0;i<len;++i){
int v = G[u][i];
if(vis[v] == 0 && v != fa)
getDeep(v,u, d);
}
}
long long calu(int u, int d){
num = 0;
getDeep(u,0,d);
long long tot = 0;
for(int i=0;i<=req;++i)
Hash[i] = 0;
for(int i=0;i<num;++i){
Hash[N[i]]++;
}
for(int i=0;i<num;++i){
for(int j=N[i];j;j = N[i] & (j-1)){
tot += Hash[j ^ req];
}
tot += Hash[req];
}
return tot;
}
void dfs(int u){
int root = 0;
int mx = INT_MAX;
int n = getSize(u,0);
getRoot(u,0,root,mx,n);
vis[root] = 1;
sum += calu(root,0);
int len = G[root].size();
for(int i = 0;i < len;++i){
int v = G[root][i];
if(vis[v] == 0){
sum -= calu(v,pre[colour[root]]);
dfs(v);
}
}
}
void init(int n) {
for (int i = 1; i <= n; ++i) {
G[i].clear();
}
sum = 0;
memset(vis,0,sizeof(vis));
}
int main(){
for(int i=0;i<11;++i) pre[i] = (1 << i);
while(scanf("%d %d",&n,&k) != EOF) {
init(n);
int u,v;
req = pre[k] - 1;
for(int i=1;i<=n;++i){
scanf("%d",colour + i);
colour[i]--;
}
for (int i = 0; i < n - 1; ++i) {
scanf("%d %d",&u,&v);
G[u].push_back(v);
G[v].push_back(u);
}
dfs(1);
printf("%lld\n",sum);
}
return 0;
}