P4084 [USACO17DEC] Barn Painting G题解——树形dp

这道题不算太难，不过由于我没开long long wa了好几发了，特此记录一下。

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

题目大意（改编自洛谷） ;给定一颗 N 个节点组成的树，你要给每个点涂上三种颜色之一，其中有 K 个节点已染色，要求任意两相邻节点颜色不同，求合法染色方案数。答案对 10e9+7 取模。其中在第一行输入N，K，分别表示前N个结点与K个已经染色的点，前N-1行输入边，后面输入b,c,表示b号结点染成c颜色。

这道题其实与P1352 没有上司的舞会这题非常相似,由于相邻的点不能染上同一颜色，所以当我们枚举当前结点为1号颜色，子节点只能染上2与3号颜色，即可状态转移方程：

其中表示u号结点涂c号颜色的方案数,假如当前结点染上1号颜色，那么就要乘上所有子结点染上2，3号颜色的方案数,由叶子结点向上递归，最后输出的结果就是，这里还有一个难点在于处理已经染上过颜色的点，这里我们只需要记录一下哪些点已经染上颜色，这样在每次递归初始化时，只用初始化已经被染上的颜色即可：

#include<set>
#include<map>
#include<cmath>
#include<stack>
#include<queue>
#include<vector>
#include<string>
#include<cstring>
#include<iomanip>
#include<iostream>
#include<algorithm>
using namespace std;
using ll=long long ;
using pii=pair<int,int> ;
#define endl "\n"

constexpr int N=2e5+10;
constexpr int mod=1e9+7;

ll dp[N][3],fla[N];
int h[N],e[N],nex[N],idx;

void add(int a,int b)
{
    e[idx]=b,nex[idx]=h[a],h[a]=idx++;
}

void dfs(int u,int fa,bool t)
{
    //没染就全部初始化
    if(!t)dp[u][0]=dp[u][1]=dp[u][2]=1;
    else 
    {
        //染过了只用初始化已经染过的点即可
        for(int i=0;i<3;i++)
        if(dp[u][i])dp[u][i]=1;
    }
    
    for(int i=h[u];i!=-1;i=nex[i])
    {
        int v=e[i];
        if(v==fa)continue;
        dfs(v,u,fla[v]);
        
        dp[u][0]=dp[u][0]*(dp[v][1]+dp[v][2])%mod;
        dp[u][1]=dp[u][1]*(dp[v][0]+dp[v][2])%mod;
        dp[u][2]=dp[u][2]*(dp[v][0]+dp[v][1])%mod;
    }
    
}

void solve()
{
    int n,k;
    cin>>n>>k;
    
    int m=n-1;
    memset(h,-1,sizeof h);
    while(m--)
    {
        int a,b;
        cin>>a>>b;
        
        add(a,b),add(b,a);
    }
    
    while(k--)
    {
        int u,c;
        cin>>u>>c;
        
        //记录已经被染上颜色的点
        fla[u]=1;
        dp[u][c-1]=1;
    }
    
    dfs(1,-1,fla[1]);
    
    cout<<(dp[1][0]+dp[1][1]+dp[1][2])%mod<<endl;
}

signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    
    int t;
    t=1;
    
    while(t--)
    solve();
    return 0;
}