HDU5977 Garden of Eden 树分治+高维前缀和

题目链接： http://acm.hdu.edu.cn/showproblem.php?pid=5977

题意：给一个有n（1<=n<=5e4）个节点的树，每个节点有一个颜色，共有k（1<=k<=7）种不同的颜色，问有多少个点对(u,v)（注意(u,v)和(v,u)算作两个答案），这些点对之间的路径上有k个不同的颜色。

题面很长，题意还是有点隐蔽的，比赛的时候刚刚敲完C题，只剩一个小时，因为不懂枚举子集和高维前缀和相关的知识（果然还是太菜了……）dp也来不及想，和队友一起码出了一个暴力的树分治，然后不断TLE，不断改，不断TLE，不断改…………就这样结束了比赛的最后一个小时。

思路：暴力枚举就超时了，但是只要知道高维前缀和以及枚举子集的相关知识，就能用O（k*(1<<k)）结束统计，外面套上树分治，O（n*logn+n*k*(1<<k) ）的时间

树分治还是那样的树分治，用状态压缩记录路径上遇到的点，接着用高维前缀和统计每个状态和他们的超集的个数，一次分治的计数方法就是num[i]*cnt[i^(1<<k-1)]，num[i]记录的是状态i的出现次数，cnt[i]是状态i和i超集的总个数，统计每个i。

高维前缀和是一种计数方法，只是短短的几行，用来计算包含i的集合的总个数。其他部分应该都ok，详见代码。

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
#define MS(x,y) memset(x,y,sizeof(x))
typedef long long LL;
const int MAXN=5e4+5;
struct Edge{
    int v,nxt;
}edge[MAXN<<1];
int head[MAXN],edgenum;
void addedge(int u,int v){
    edge[edgenum].v=v;
    edge[edgenum].nxt=head[u];
    head[u]=edgenum++;
}
int n,k,K;
int col[MAXN];
LL ans;

int root,max_sub;
int siz[MAXN];
LL cnt[MAXN],cnt1[MAXN];
bool used[MAXN];
void findroot(int u,int f,int num){
    siz[u]=1;
    int K=0;
    for(int i=head[u];~i;i=edge[i].nxt){
        int &v=edge[i].v;
        if(used[v]||v==f) continue;
        findroot(v,u,num);
        K=max(K,siz[v]);
        siz[u]+=siz[v];
    }
    K=max(K,num-siz[u]);
    if(K<max_sub) max_sub=K,root=u;
}
void dfs(int u,int f,int sta){
    sta|=(1<<col[u]);
    ++cnt[sta];
    ++cnt1[sta];
    for(int i=head[u];~i;i=edge[i].nxt){
        int &v=edge[i].v;
        if(used[v]||v==f) continue;
        dfs(v,u,sta);
    }
}
LL work2(int u,int sta){
    for(int i=0;i<=K;++i) cnt[i]=cnt1[i]=0;
    sta|=(1<<col[u]);
    ++cnt[sta];
    ++cnt1[sta];
    for(int i=head[u];~i;i=edge[i].nxt){
        int &v=edge[i].v;
        if(used[v]) continue;
        dfs(v,u,sta);
    }
    for(int i=0;i<k;++i) for(int j=K;j>=0;--j){ //高维前缀和
        if(!((1<<i)&j)) cnt[j]+=cnt[(1<<i)|j];
    }
    LL ret=0;
    for(int i=1;i<=K;++i) ret+=cnt1[i]*cnt[(K^i)];
    return ret;
}
LL work1(int u,int num){
    max_sub=num;
    root=u;
    findroot(u,0,num);
    u=root;
    used[u]=true;
    LL ret=work2(u,0);
    for(int i=head[u];~i;i=edge[i].nxt){
        int &v=edge[i].v;
        if(used[v]) continue;
        ret-=work2(v,(1<<col[u]));
        ret+=work1(v,siz[v]);
    }
    return ret;
}

int main(){//freopen("in.txt","r",stdin);
    while(~scanf("%d%d",&n,&k)){
        ans=0;
        K=(1<<k)-1;
        MS(head,-1);
        MS(used,false);
        edgenum=0;
        for(int i=1;i<=n;++i) scanf("%d",col+i),--col[i];
        int u,v;
        for(int i=1;i<n;++i){
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        printf("%I64d\n",work1(1,n));
    }
}