【HDOJ5977】Garden of Eden（点分治）

题意：给定一棵n个点的树，每个节点上有一种颜色a[i]，一共有k种颜色，问包含所有颜色的路径条数

n<=5e4，k<=10

思路：点分治求方案数

集合并卷积的时候暴力枚举状态即可O(n^logn*2^k)

75e的复杂度 只跑了1.7s 我也是醉了

 #include<cstdio>
 #include<cstring>
 #include<iostream>
 #include<algorithm>
 #include<cmath>
 typedef long long ll;
 using namespace std;
 #define N   110000
 #define oo  10000000
 #define MOD 1000000007

 ll ans;
 int head[N],vet[N],nxt[N],a[N],flag[N],son[N],f[N],s[],dep[N],d[N],
     n,k,all,tot,sum,root;

 void add(int a,int b)
 {
     nxt[++tot]=head[a];
     vet[tot]=b;
     head[a]=tot;
 }

 void getroot(int u,int fa)
 {
     son[u]=; f[u]=;
     int e=head[u];
     while(e)
     {
         int v=vet[e];
         if(v!=fa&&!flag[v])
         {
             getroot(v,u);
             son[u]+=son[v];
             f[u]=max(f[u],son[v]);
         }
         e=nxt[e];
     }
     f[u]=max(f[u],sum-f[u]);
     if(f[u]<f[root]) root=u;
 }

 void getdep(int u,int fa)
 {
     dep[++dep[]]=d[u];
     int e=head[u];
     while(e)
     {
         int v=vet[e];
         if(v!=fa&&!flag[v])
         {
             d[v]=d[u]|(<<a[v]);
             getdep(v,u);
         }
         e=nxt[e];
     }
 }

 ll calc(int u,int now)
 {
     d[u]=now; dep[]=;
     getdep(u,);
     memset(s,,sizeof(s));
     ll ans=;
     for(int i=;i<=dep[];i++)
     {
          for(int j=;j<=all;j++)
            if((dep[i]|j)==all) ans+=s[j];
           s[dep[i]]++;
     }
     return ans;
 }

 void solve(int u)
 {
     ans+=calc(u,<<a[u]);
     flag[u]=;
     int e=head[u];
     while(e)
     {
         int v=vet[e];
         if(!flag[v])
         {
             ans-=calc(v,(<<a[u])|(<<a[v]));
             sum=son[v];
             root=;
             getroot(v,root);
             solve(root);
         }
         e=nxt[e];
     }
 }

 int main()
 {
     while(scanf("%d%d",&n,&k)!=EOF)
     {
         all=(<<k)-;
         for(int i=;i<=n;i++) head[i]=flag[i]=;
         for(int i=;i<=n;i++)
         {
             scanf("%d",&a[i]);
             a[i]--;
         }
         tot=;
         for(int i=;i<=n-;i++)
         {
             int x,y;
             scanf("%d%d",&x,&y);
             add(x,y);
             add(y,x);
         }
         if(k==)
         {
             ll ans=n*n;
             printf("%I64d\n",ans);
             continue;
         }
         sum=n; f[]=oo; ans=; root=;
         getroot(,);
         solve(root);
         ans*=;
         printf("%I64d\n",ans);
     }
     return ;
 }