洛谷 P3914 染色计数

P3914 染色计数

题目描述

有一颗NN个节点的树，节点用1,2,\cdots,N1,2,⋯,N编号。你要给它染色，使得相邻节点的颜色不同。有MM种颜色，用1,2,\cdots,M1,2,⋯,M编号。每个节点可以染MM种颜色中的若干种，求不同染色方案的数量除以(10^9 + 7109+7)的余数。

输入输出格式

输入格式：

 

第1 行，2 个整数N,MN,M。

接下来NN行，第ii行表示节点ii可以染的颜色。第1个整数k_iki​，表示可以染的颜色数量。接下来k_iki​个整数，表示可以染的颜色编号。

最后N - 1N−1行，每行2个整数A_i,B_iAi​,Bi​，表示边(A_i,B_i)(Ai​,Bi​)。

 

输出格式：

 

1 个整数，表示所有的数。

 

输入输出样例

输入样例#1： 复制

2 2
1 1
2 1 2
1 2

输出样例#1： 复制

1

说明

• 对于30% 的数据，1 \le N \le 10; 1 \le M \le 41≤N≤10;1≤M≤4；

• 对于60% 的数据，1 \le N \le 200; 1 \le M \le 2001≤N≤200;1≤M≤200；

• 对于100% 的数据，1 \le N \le 5000; 1 \le M \le 50001≤N≤5000;1≤M≤5000。

思路：组合数学+树形DP。

f[i][j]表示在以i为根的子树中，当i号节点的颜色为j时的方案数。

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define mod 1000000007
#define MAXN 5010
using namespace std;
int n,m,tot;
int ans[MAXN],f[5010][5010];
int to[MAXN*2],net[MAXN*2],head[MAXN*2];
void add(int u,int v){
    to[++tot]=v;net[tot]=head[u];head[u]=tot;
    to[++tot]=u;net[tot]=head[v];head[v]=tot;
}
void dfs(int now,int fa){
    for(int i=head[now];i;i=net[i])
        if(to[i]!=fa)
            dfs(to[i],now);
    for(int i=1;i<=m;i++)
        if(f[now][i]){
            for(int j=head[now];j;j=net[j])
                if(to[j]!=fa)
                    f[now][i]=1LL*f[now][i]*(ans[to[j]]-f[to[j]][i])%mod;
            while(f[now][i]<0)    f[now][i]+=mod;
            ans[now]=(ans[now]+f[now][i])%mod;
        }
}
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++){
        int k;scanf("%d",&k);
        for(int j=1;j<=k;j++){
            int q;scanf("%d",&q);
            f[i][q]++;
        }
    }
    for(int i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
    }
    add(0,1);
    dfs(0,0);
    printf("%d",ans[1]);
}