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));
}
}