P3914 染色计数【树形DP】

题目描述
有一颗N个节点的树，节点用1,2,⋯,N编号。你要给它染色，使得相邻节点的颜色不同。有M种颜色，用1,2,⋯,M编号。每个节点可以染M种颜色中的若干种，求不同染色方案的数量除以(10^9 + 7)的余数。
输入输出格式
输入格式：
第1 行，2 个整数N,M。 1≤N≤5000;1≤M≤5000
接下来N行，第i行表示节点ii可以染的颜色。第1个整数k，表示可以染的颜色数量。接下来k个整数，表示可以染的颜色编号。
最后N - 1行，每行2个整数A_i,B_i，表示边(A_i,B_i)
输入样例#1：
2 2
1 1
2 1 2
1 2
输出样例#1：
1

分析：【树形DP】
这是一道很经典的树形DP的题目，相邻点染的颜色不同，对于以i点作为根节点的子树，你要记录当前i节点用第j类颜色的方案数$dp[i][j]$，以及用m种颜色的总方案数$p[i]$。转移的时候，$dp[i][j]=dp[i][j]\ast (pp[son]-dp[son][j]+mod)$，$p[i]={\sum }_{j=1}^{m}dp[i][j]$.

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <cmath>
#include <vector>
using namespace std;
typedef long long LL;
const LL mod = 1e9 + 7;
const int INF = 0x3f3f3f3f;
const int MAXN = 5010;
vector<int> G[MAXN];
int dp[MAXN][MAXN], pp[MAXN];
bool vis[MAXN][MAXN];
int n, m, k, col;

void dfs(int x, int fa) {
    int u, v;
    if(fa > 0 && G[x].size() == 1) {
        int cnt = 0;
        for(int i = 1; i <= m; ++i) {
            if(!vis[x][i]) continue;
            dp[x][i] = 1;
            cnt++;
        }
        pp[x] = cnt;
        return ;
    }
    for(int i = 0; i < G[x].size(); ++i) {
        u = G[x][i];
        if(u == fa) continue;
        dfs(u, x);
        //pp[x] = (pp[x] * 1ll * pp[u]) % mod; //这样乘出来答案是不合法的
    }
    for(int i = 1; i <= m; ++i) {
        if(!vis[x][i]) continue;
        dp[x][i] = 1;
        for(int j = 0; j < G[x].size(); ++j) {
            v = G[x][j];
            if(v == fa) continue;
            dp[x][i] = (dp[x][i] * 1ll * ((pp[v] - dp[v][i] + mod) % mod)) % mod;
           // printf("## %d %d %d %d\n", dp[x][u], pp[v], x, u);
            //printf("$$ %d %d %d %d\n", dp[v][u], pp[v], v, u);
        }
        pp[x] = (1ll * pp[x] + 1ll * dp[x][i]) % mod; //累加计和
    }
}

int main() {
    scanf("%d %d", &n, &m);
    for(int i = 1; i <= n; ++i) {
        scanf("%d", &k);
        for(int j = 1; j <= k; ++j) {
            scanf("%d", &col);
            vis[i][col] = true;
        }
    }
    for(int i = 1; i < n; ++i) {
        int a, b;
        scanf("%d %d", &a, &b);
        G[a].push_back(b);
        G[b].push_back(a);
    }
    dfs(1, -1); int sum = 0;
    for(int i = 1; i <= m; ++i) {
        if(!vis[1][i]) continue;
        //printf("# %d\n", dp[1][i]);
        sum += dp[1][i]; sum %= mod;
    }
    printf("%d\n", sum);
    return 0;
}

/*
2 3
2 1 2
2 1 3
1 2


4 3
3 1 2 3
3 1 2 3
2 1 2
2 1 3
1 4
4 2
4 3

*/