HDU 5977 Garden of Eden （点分治+子集枚举）

Posted 2021-01-12 lalalatianlalu

题意：给你一颗树，每个节点上都有一个种类，问你包含所有种类的路径有多少条（1,3）（3,1）算2条

思路：最近狂补点分治，这道题点分治的部分还是很好理解的，直接套点分治的板子，只不过这道题不需要使用dep数组，我们知道dep其实维护的是从rt到子节点的权值，这里因为只有10种种类，所以我们直接状压（k<10，明示状压）,然后就是子集枚举加速遍历（好久没写过这个了，记得之前一次写是在一次状压dp里，里面用到了子集枚举加速，但是题目已经记不清楚了），在这里复习了一下，（附一篇网上的题解，传送门）

代码：（写完点分治忘记第一次找根的时候要对f[0]=sum=n，这条赋值，找了很久）

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long LL;

const int maxn=5e4+7;
struct Edge
{
    int to,next;
}a[maxn<<1];
int n,k,head[maxn],cnt;
int root,sum,vis[maxn],sz[maxn];
int f[maxn],stk[maxn],top;
LL ans;
LL Hash[1200];
int col[maxn];
void addedge(int u,int v)
{
    a[++cnt].to=v;
    a[cnt].next=head[u];
    head[u]=cnt;
}
void init()
{
    cnt=0;
    memset(head,0,sizeof(head));
    for(int i=0;i<=n+3;++i)vis[i]=0;
    ans=0;root=0;
}
void getroot(int u,int fa)
{
    sz[u]=1;f[u]=0;
    for(int e=head[u];e;e=a[e].next){
        int v=a[e].to;
        if(v==fa||vis[v])continue;
        getroot(v,u);
        sz[u]+=sz[v];
        f[u]=max(f[u],sz[v]);
    }
    f[u]=max(f[u],sum-sz[u]);
    if(f[u]<f[root])root=u;
}

void getdep(int u,int fa,int s)
{
    stk[++top]=s;
    for(int e=head[u];e;e=a[e].next){
        int v=a[e].to;
        if(v==fa||vis[v])continue;
        getdep(v,u,s|(1<<col[v]));
    }
}
LL calc(int u,int s)
{
    top=0;
    getdep(u,0,s);
    LL res=0;
    memset(Hash,0,sizeof(Hash));
    for(int i=1;i<=top;i++)Hash[stk[i]]++;
    for(int i=1;i<=top;i++){
        Hash[stk[i]]--;
        res+=Hash[(1<<k)-1];
        for(int s0=stk[i];s0;s0=(s0-1)&stk[i]){
            res+=Hash[((1<<k)-1)^s0];
        }
        Hash[stk[i]]++;
    }
    return res;
}

void solve(int u)
{
    ans+=calc(u,(1<<col[u]));
    vis[u]=1;
    for(int e=head[u];e;e=a[e].next){
        int v=a[e].to;
        if(vis[v])continue;
        ans-=calc(v,(1<<col[u])|(1<<col[v]));
        sum=sz[v];root=0;
        getroot(v,0);
        solve(root);
    }
}

int main()
{
    while(~scanf("%d%d",&n,&k)){
        init();
        for(int i=1;i<=n;++i){
            scanf("%d",&col[i]);
            col[i]--;
        }
        for(int i=1;i<n;++i){
            int u,v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        if(k==1){
            printf("%d
",n*n);
            continue;
        }
        ans=0;
        sum=f[0]=n;
        getroot(1,0);
        solve(root);
        printf("%lld
",ans);
    }
    return 0;
}