【树形计数DP】P4084 [USACO17DEC]Barn Painting G-蒲公英云

【树形计数DP】P4084 [USACO17DEC]Barn Painting G

P4084 [USACO17DEC]Barn Painting G - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

题意：

思路：

简单树形DP题，应该要会的

树上染色，且颜色一共就3种，很自然地想到这样设计：

dp[u][j]表示以u为根的子树，结点u涂成颜色j的方案数

转移的时候容易想不清楚，其实就是乘法原理，想成一列，然后类似于二分图连边的感觉，就行了

Code：

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int mxn=1e5+10;
const int mxe=1e5+10;
const int mod=1e9+7;
struct ty{
    int to,next;
}edge[mxe<<2];
int N,K,x,c;
int u,v,tot=0;
int head[mxn];
int dp[mxn][4],col[mxn];
void add(int u,int v){
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void G_init(){
    tot=0;
    for(int i=0;i<=N;i++){
        head[i]=-1;
    }
}
void dfs(int u,int fa){
    if(!col[u]) dp[u][1]=dp[u][2]=dp[u][3]=1;
    else dp[u][col[u]]=1;
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        dfs(edge[i].to,u);
        dp[u][1]*=(dp[edge[i].to][2]+dp[edge[i].to][3])%mod,dp[u][1]%=mod;
        dp[u][2]*=(dp[edge[i].to][1]+dp[edge[i].to][3])%mod,dp[u][2]%=mod;
        dp[u][3]*=(dp[edge[i].to][1]+dp[edge[i].to][2])%mod,dp[u][3]%=mod;
    }
}
void solve(){
    cin>>N>>K;
    G_init();
    for(int i=1;i<=N-1;i++){
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=K;i++){
        cin>>x>>c;
        col[x]=c;
    }
    dfs(1,0);
    cout<<(dp[1][1]+dp[1][2]+dp[1][3])%mod<<'\n';
}
signed main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    int __=1;//cin>>__;
    while(__--)solve();return 0;
}