[USACO17DEC]Barn Painting (树形$dp$)

题目链接

Solution

比较简单的树形 (dp) 。
(f[i][j]) 代表 (i) 为根的子树 ，(i) 涂 (j) 号颜色的方案数。
转移很显然 :

[f[i][1]=prod(f[t][2]+f[t][3]) ]

其中 (k) 代表它的子节点。 其他两种颜色以此类推。
但需要注意的是对于颜色固定的点，除固定的颜色外，其他两种颜色的方案要赋为 (0) 。
PS : 要开 longlong ，以及还要 mod 1000000007 。

Code

#include<bits/stdc++.h>
#define N 100010
#define ll long long
#define inf 0x3f3f3f3f
#define mod 1000000007
using namespace std;

void in(ll &x)
{
    char ch=getchar();ll f=1,w=0;
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){w=w*10+ch-'0';ch=getchar();}
    x=f*w; return;
}

struct sj{
    ll to,next;
}a[N*2];
ll head[N],size;
ll n,col[N],k;

void add(ll x,ll y)
{
    a[++size].to=y;
    a[size].next=head[x];
    head[x]=size;
}
ll f[N][4];

void dfs(ll x,ll fr)
{
    f[x][1]=f[x][2]=f[x][3]=1;
    for(ll i=head[x];i;i=a[i].next)
    {
        ll tt=a[i].to;
        if(tt==fr)continue;
        dfs(tt,x);
        f[x][1]*=(f[tt][2]+f[tt][3])%mod;
        f[x][1]%=mod;
        f[x][2]*=(f[tt][1]+f[tt][3])%mod;
        f[x][2]%=mod;
        f[x][3]*=(f[tt][2]+f[tt][1])%mod;
        f[x][3]%=mod;
    }
    if(col[x])
    for(int i=1;i<=3;i++)
    if(col[x]!=i)f[x][i]=0;
    return;
}

int main()
{
    in(n); in(k);
    for(ll i=1;i<n;i++)
    {
        ll x,y; in(x),in(y);
        add(x,y),add(y,x);
    }
    for(ll i=1;i<=k;i++)
    {
        ll x,y; in(x),in(y);
        col[x]=y; 
    }
    dfs(1,0);
    cout<<(f[1][1]%mod+f[1][2]%mod+f[1][3]%mod)%mod<<endl;
    return 0;
}