[USACO17DEC] Barn Painting - 树形dp

设\(f[i][j]\)为\(i\)子树，当\(i\)为\(j\)时的方案数

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1000005;
const int mod = 1e+9+7;
vector <int> g[N];
int vis[N],n,m,t1,t2,f[N][4],c[N];

void dfs(int p) {
    vis[p]=1;
    int flag = 0;
    f[p][1]=f[p][2]=f[p][3]=1;
    for(int i=0;i<g[p].size();i++) {
        int q=g[p][i];
        if(vis[q]) continue;
        flag = 1;
        dfs(q);
        for(int j=1;j<=3;j++) {
            int tmp=0;
            for(int k=1;k<=3;k++) {
                if(j!=k) tmp+=f[q][k];
            }
            f[p][j]*=tmp, f[p][j]%=mod;
        }
    }
    if(c[p]) {
        for(int i=1;i<=3;i++) {
            if(i==c[p]) continue;
            f[p][i]=0;
        }
    }
}

signed main() {
    scanf("%lld%lld",&n,&m);
    for(int i=1;i<n;i++) {
        scanf("%lld%lld",&t1,&t2);
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }
    for(int i=1;i<=m;i++) {
        scanf("%lld%lld",&t1,&t2);
        c[t1]=t2;
    }
    dfs(1);
    cout<<(f[1][1]+f[1][2]+f[1][3])%mod<<endl;
}