Hackrank Kingdom Division 树形DP

题目链接：传送门

题意：

　　给你一棵树，n个点

　　每个点可以染成红色和蓝色

　　但是红色的点与其相邻的点中必须有红色节点，蓝色也是

　　问你有多少种染色的方案

题解：

　　树形dp

　　先转化为有根树，取1为根

　　设定dp[now][red][red] 表示的是当前now节点然red色，其父亲节点染red色的可行方案数

　　转移很容易想到

#include<bits/stdc++.h>
using namespace std;
#pragma comment(linker, "/STACK:102400000,102400000")
#define ls i<<1
#define rs ls | 1
#define mid ((ll+rr)>>1)
#define pii pair<int,int>
#define MP make_pair
typedef long long LL;
const long long INF = 1e18+1LL;
const double Pi = acos(-1.0);
const int N = 5e5+10, M = 1e2+20,inf = 2e9;
const LL mod = 1e9+7;

LL dp[N][2][2];
vector<int > G[N];
int n;
void dfs(int u,int f) {
    for(int i = 0; i < G[u].size(); ++i) {
        int to = G[u][i];
        if(to == f) continue;
        dfs(to,u);
    }
    dp[u][1][1] = 1;
    dp[u][0][0] = 1;
    int ok = 0;
    LL ans1=  1,ans2 = 1;
    for(int i = 0; i < G[u].size(); ++i) {
        int to = G[u][i];
        if(to == f) continue;
        ok = 1;
        dp[u][1][1] *= (dp[to][0][1]+dp[to][1][1])%mod;
        dp[u][1][1] %= mod;

        dp[u][0][0] *= (dp[to][0][0]+dp[to][1][0])%mod;
        dp[u][0][0] %= mod;
        ans1 *= (dp[to][0][1]) % mod;
        ans1 %= mod;
        ans2 *= (dp[to][1][0]) % mod;
        ans2 %= mod;
    }
    if(ok) {
        dp[u][1][0] = (dp[u][1][1] - ans1 + mod) % mod;
        dp[u][0][1] = (dp[u][0][0] - ans2 + mod) % mod;
    }
}

int main() {
    scanf("%d",&n);
    for(int i = 1; i < n; ++i) {
        int x,y;
        scanf("%d%d",&x,&y);
        G[x].push_back(y);
        G[y].push_back(x);
    }
    dfs(1,0);
    printf("%lld
",(dp[1][0][1]+dp[1][1][0]) % mod);
    return 0;
}