2017 Wuhan University Programming Contest (Online Round) B Color 树形dp求染色方法数

/**
题目：Color
链接：https://oj.ejq.me/problem/23
题意：给定一颗树，将树上的点最多染成m种颜色，有些节点不可以染成某些颜色。相邻节点颜色不同。求染色方法数。
思路：树形dp，定义dp[i][j]表示以i为根，i节点染色为j时候的子树的染色方法数。

*/
#include<bits/stdc++.h>
#define LL long long
using namespace std;
typedef long long ll;
const int maxn = 2e4+100;
ll dp[maxn][22];///dp[i][j]表示以i为根，i节点染色为j时候的子树的染色方法数。
ll mod = 1e7+9;
vector<int> G[maxn];///图
vector<int> color[maxn];///节点i可以染的颜色；
ll cnt[maxn];///记录节点i的子树的方法数。可以减少一层循环。优化时间。
int n, m;
void init()
{
    for(int i = 1; i <= n; i++){
        G[i].clear();
        color[i].clear();
    }
    memset(cnt, 0, sizeof cnt);
    memset(dp, 0, sizeof dp);
}
void dfs(int r,int f)
{
    int lenr = G[r].size();
    int lenc = color[r].size();
    for(int j = 0; j < lenr; j++){
        int &v = G[r][j];
        if(v==f) continue;
        dfs(v,r);
    }
    for(int i = 0; i < lenc; i++){
        ll res = 1;
        int &c = color[r][i];
        for(int j = 0; j < lenr; j++){
            int &v = G[r][j];
            if(v==f) continue;
            res = res*(cnt[v]-dp[v][c]+mod)%mod;
        }
        dp[r][c] = res;
        cnt[r] = (cnt[r]+dp[r][c])%mod;
    }
}
int main()
{
    while(scanf("%d%d",&n,&m)==2)
    {
        init();
        int x, y;
        for(int i = 1; i < n; i++){
            scanf("%d%d",&x,&y);
            G[x].push_back(y);
            G[y].push_back(x);
        }
        for(int i = 1; i <= n; i++){
            for(int j = 1; j <= m; j++){
                scanf("%d",&x);
                if(x){
                    color[i].push_back(j);
                }
            }
        }

        dfs(1,1);

        printf("%lld
",cnt[1]);
 /*       for(int i = 1; i <= n; i++){
            for(int j = 1; j <= m; j++){
                printf("%d ",dp[i][j]);
            }
            cout<<endl;
        }*/

    }
    return 0;
}