【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[1100],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]=1; f[u]=0;
    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[0]]=d[u];
    int e=head[u];
    while(e)
    {
        int v=vet[e];
        if(v!=fa&&!flag[v])
        {
            d[v]=d[u]|(1<<a[v]);
            getdep(v,u);
        }
        e=nxt[e];
    }
}

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

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

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