Black and Red vertices of Tree CodeChef - BLREDSET (树型dp套路题，可惜俺之前8会)

题目大意

给定一棵N个结点的树，树上有三种颜色，0，1，2分别代表白，黑，红。
问题是有多少个联通块都为白色，且删掉这个联通块至少有一个红黑结点不相连。

思路

树上有一些关键点为白色且至少有两颗子树分别含有红黑两种结点，我们只需统计出这样的结点有哪些，然后求出其至少包含一个关键点的联通图有几种就行了，但是求”至少包含一个“不好求说实话，可以求出所有以r为根的白色联通图个数，减去以r为根不含有关键点的联通图个数即可。

求以u为根的联通图个数可以用累乘的方法，分别然后乘上u其它子树大小 + 1。

代码

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
using namespace std;

typedef long long ll;
const ll mod = 1e9 + 7;
const int maxn = 1e5 + 100;
int n;
vector<int> G[maxn];
int sz[maxn][3],col[maxn];
bool good[maxn];
ll tot[maxn], cntbad[maxn];
void dfs0(int u, int fa) {
    for(int i = 0; i < 3; i++) sz[u][i]=0;
    sz[u][col[u]] = 1;
    bool havered=false, haveblack=false;
    for(auto & v : G[u]) {
        if(v == fa) continue;
        dfs0(v,u);
        
        good[u] |= (havered && sz[v][1]);
        good[u] |= (haveblack && sz[v][2]);
        haveblack |= sz[v][1];
        havered |= sz[v][2];
        
        for(int i = 0; i < 3; i++) sz[u][i]+=sz[v][i];
    }
}
void dfs1(int u, int fa) {
    tot[u] = !col[u];
    cntbad[u] = (!col[u] && !good[u]);
    for(auto & v : G[u]) {
        if(v == fa) continue;
        dfs1(v, u);
        tot[u] = tot[u] * (tot[v] + 1) % mod;
        cntbad[u] = cntbad[u] * (cntbad[v] + 1) % mod;
    }
}
int main() {
    //freopen("/Users/maoxiangsun/MyRepertory/input.txt", "r", stdin);
    int T;
    cin >> T;
    while(T--) {
        cin>>n;
        for(int i = 1; i <= n; i++) {good[i]=false;G[i].clear();}
        for(int i = 0; i < n - 1; i++) {
            int u,v;cin>>u>>v;
            G[u].push_back(v);
            G[v].push_back(u);
        }
        for(int i = 1; i <= n; i++) cin>>col[i];
        
        dfs0(1, -1);
        for(int u = 1; u <= n; u++) {
            good[u] |= (sz[u][2] && sz[1][1] - sz[u][1]);
            good[u] |= (sz[u][1] && sz[1][2] - sz[u][2]);
        }
        dfs1(1, -1);
        ll ans = 0;
        for(int i = 1; i <= n; i++) {
            ans = (ans + tot[i] - cntbad[i] + mod) % mod;
        }
        printf("%lld\n",ans);
    }
    return 0;
}
/*
 2
 6
 1 2
 1 3
 1 4
 3 5
 3 6
 0 1 0 1 2 0
 6
 1 2
 1 3
 1 4
 3 5
 3 6
 1 0 0 0 2 0
 */