洛谷P4084 [USACO17DEC]Barn Painting G 题解

题目链接：P4084 [USACO17DEC]Barn Painting G

题意：题意：给定一颗N个节点组成的树，3种颜色，其中K个节点已染色，要求任意两相邻节点颜色不同，求合法染色方案数。

这个题太水了

设 \(dp[u][1/2/3]\) 表示 \(u\) 结点涂 \(1/2/3\) 颜色的方案数

根据加法原理，

如果 \(u\) 已经有颜色了，那么 \[ dp[u][j] = \sum dp[v][k]\times [k\ne j]\times [j=\text{val[u]}] \] 否则 \[ dp[u][j] = \sum dp[v][k]\times [k\ne j] \] 时间复杂度 \(O(n)\)

代码：

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
#define N (int)(1e5+15)
const int p=1e9+7;
int n,k;
int val[N],dp[N][4];
vector<int> vec[N];
void add(int &a,int b){a=(a+b%p)%p;}
void mul(int &a,int b){a=a%p*b%p;}
void dfs(int u,int f)
{
    for(int v : vec[u])
    {
        if(v==f)continue;
        dfs(v,u);
        for(int i=1; i<=3; i++)
        {
            int res=0;
            for(int j=1; j<=3; j++)
                if(i!=j)add(res,dp[v][j]);
            mul(dp[u][i],res);
        }
    }
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n >> k;
    for(int i=1,u,v; i<n; i++)
    {
        cin >> u >> v;
        vec[u].push_back(v);
        vec[v].push_back(u);
    }
    for(int i=1; i<=n; i++)
        for(int j=1; j<=3; j++)
            dp[i][j]=1;
    for(int i=1,u; i<=k; i++)
    {
        cin >> u >> val[u];
        for(int j=1; j<=3; j++)
            if(val[u]!=j)dp[u][j]=0;
    }
    dfs(1,1);
    cout << ((dp[1][1]+dp[1][2]+dp[1][3])%p+p)%p << endl;
    return 0;
}