HDU-5977 树的分治+子集枚举

题目

Garden of Eden

题意

求树上满足条件的点对的个数，条件是在这两点间不经过重复的点能吃到的苹果类型的总数大于 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;
}