[USACO17DEC] Barn Painting G（树形dp）

green qwq

于 2024-05-28 10:00:00 发布

阅读量1.3k

点赞数 40

文章标签：图论深度优先算法

本文链接：https://blog.csdn.net/m0_60738889/article/details/139203006

版权

[USACO17DEC] Barn Painting G

题面翻译

题意：给定一颗 $N$ 个节点组成的树，你要给每个点涂上三种颜色之一，其中有 $K$ 个节点已染色，要求任意两相邻节点颜色不同，求合法染色方案数。答案对 $10^9+7$ 取模。

题目描述

Farmer John has a large farm with $N$ barns ( $\le N \le 10^5$ ), some of which are already painted and some not yet painted. Farmer John wants to paint these remaining barns so that all the barns are painted, but he only has three paint colors available. Moreover, his prize cow Bessie becomes confused if two barns that are directly reachable from one another are the same color, so he wants to make sure this situation does not happen.

It is guaranteed that the connections between the $N$ barns do not form any ‘cycles’. That is, between any two barns, there is at most one sequence of connections that will lead from one to the other.

How many ways can Farmer John paint the remaining yet-uncolored barns?

输入格式

The first line contains two integers $N$ and $K$ ( $\le K \le N$ ), respectively the number of barns on the farm and the number of barns that have already been painted.

The next $N - 1$ lines each contain two integers $x$ and $y$ ( $\le x, y \le N, x \neq y$ ) describing a path directly connecting barns $x$ and $y$ .

The next $K$ lines each contain two integers $b$ and $c$ ( $\le b \le N$ , $\le c \le 3$ ) indicating that barn $b$ is painted with color $c$ .

输出格式

Compute the number of valid ways to paint the remaining barns, modulo $10^9 + 7$ , such that no two barns which are directly connected are the same color.

样例 #1

样例输入 #1

样例输出 #1

思路

这道题不就是跟那个没有上司的舞会那道题类似吗？只不过这边有3中状态。具体的状态推到可以看：
在这里插入图片描述
图片来源，讲的很详细

代码

//本道题跟上司的舞会很像，3种颜色

#include<iostream>
#include<algorithm>
#include<cstring>

#define int long long

using namespace std;

const int N = 2e5+10,mod = 1e9+7,M = 2*N;

int f[N][4];//j表示3种颜色
int e[M],ne[M],w[M],h[N],idx;
bool st[N];
int color[N];
int n,k;

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

void dfs(int u,int fa){
    if(color[u]){
        //如果已经上过色了，其他两种颜色的方案数为0。
        f[u][color[u]]=1;
    }else{
        f[u][1]=f[u][2]=f[u][3]=1;
    }
    
    for(int i=h[u];~i;i=ne[i]){
        int j=e[i];
        
        if(j==fa)continue;
        
        dfs(j,u);
        //因为这个是分步进行的，所以要乘法
        f[u][1]=f[u][1]*(f[j][2]+f[j][3])%mod;
        f[u][2]=f[u][2]*(f[j][1]+f[j][3])%mod;
        f[u][3]=f[u][3]*(f[j][1]+f[j][2])%mod;
    }
}

signed main(){
    cin>>n>>k;
    
    memset(h,-1,sizeof h);
    
    for(int i=1;i<n;i++){
        int a,b;
        cin>>a>>b;
        add(a,b);
        add(b,a);
    }
    
    for(int i=1;i<=k;i++){
        int x,y;
        cin>>x>>y;
        color[x]=y;
    }
    
    dfs(1,-1);
    
    cout<<(f[1][1]+f[1][2]+f[1][3])%mod;
    
    return 0;
}