[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;
}