【jzoj5050】【颜色树】【容斥原理】【树形动态规划】

题目大意

传说中，这棵不见边际的树有N个节点，每个节点都有1片叶子，每片叶子都拥有K种颜色中的一种，独孤玉溪喜欢爬到这棵树上，沿着一条路线摘叶子，并拥有所有颜色的叶子。独孤玉溪会选择一个起点，并沿着树边走，然后最终停在一个终点上(起点和终点可能相同)，当然了每一个结点只能经过一次(每一片叶子只能摘一遍)。独孤玉溪突生奇想，有多少种不同的方案能满足自己呢？(两种方案不同当且仅当起点不同或终点不同)。

解题思路

容斥那些颜色不选，其他可选可不选，使用树形dp统计答案。

code

#include<set>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LD double
#define LL long long
#define ULL unsigned long long
#define min(a,b) ((a<b)?a:b)
#define max(a,b) ((a>b)?a:b)
#define fo(i,j,k) for(int i=j;i<=k;i++)
#define fd(i,j,k) for(int i=j;i>=k;i--)
#define fr(i,j) for(int i=begin[j];i;i=next[i])
using namespace std;
int const mn=5*1e4+9,mm=1e5+9,inf=1e9;
int n,K,gra,begin[mn],to[mm],next[mm],f[mn],vis[mn],col[mn],tag[20];LL ans;
void insert(int u,int v){
    to[++gra]=v;
    next[gra]=begin[u];
    begin[u]=gra;
}
void dfs(int now,int tmp){
    vis[now]=1;f[now]=1;LL tmp2=0;
    fr(i,now)if((!vis[to[i]])&&(!tag[col[to[i]]])){
        dfs(to[i],tmp);
        f[now]+=f[to[i]];
        tmp2+=1ll*f[to[i]]*f[to[i]];
    }
    ans+=(1ll*f[now]*f[now]-tmp2)*tmp;
}
int main(){
    freopen("colortree.in","r",stdin);
    freopen("colortree.out","w",stdout);
    scanf("%d%d",&n,&K);
    fo(i,1,n)scanf("%d",&col[i]);int u,v;
    fo(i,1,n-1){
        scanf("%d%d",&u,&v);
        insert(u,v);insert(v,u);
    }
    int mx=(1<<K)-1;
    fo(s,0,mx){
        int S=s,tmp=0;
        fo(i,1,K){
            tag[i]=S&1;
            tmp+=tag[i];
            S>>=1;
        }
        tmp=(tmp&1)?-1:1;
        fo(i,1,n)vis[i]=0;
        fo(i,1,n)if((!vis[i])&&(!tag[col[i]]))dfs(i,tmp);
    }
    printf("%lld",ans);
    return 0;
}