BLREDSET, C++ (gcc)

Language: Layout:

#include <bits/stdc++.h>

using namespace std;

#define pb push_back
typedef long long int ll;

const int N = 1e5 + 5;

/***********************************************************
sub[i] == total number of ways we can select a subset from the sub-tree rooted at i (always include i) 
    
    used because when we need to calculate the answer for the parent of the current node we have to take care of the condition "the unclored components needs to be connnected always"
    so, include 'i' so that for parent this will take care of it's subtree.
    
    
red[i] == number of red nodes in the sub-tree rooted at i.

avail[i] == number of valid ways to choose a set of uncolored nodes from the sub-tree rooted at 'i'.

**************************************************************/

vector<ll> g[N], sub(N, 0), black(N, 0), red(N, 0), color(N, 0), avail(N, 0);

const int mod = 1e9 + 7;

ll n = 0, ans = 0, tot_r, tot_b;

void dfs(int s, int p){
    ll &b = black[s], &r = red[s], &sb = sub[s];
    sb = 1;
    int fg = 0;
    ll x = 1;
    
    for(auto i : g[s]){
        if(i == p) continue;
        
        dfs(i, s);
        if(black[i] && r) fg = 1;
        if(red[i] && b) fg = 1;
        
        b += black[i];
        r += red[i];
        sb = (sb * (sub[i] + 1)) % mod;               // there are sub[i] ways to select a sub-tree from the tree rooted at i such that they contains only unclolored nodes
                                                     // so there are two possible cases : either we take a set from those sets or we leave that sub-tree blank i.e. don't take anyhing from there.
        
        x = (x * (sub[i] - avail[i] + mod + 1) % mod) % mod;   // now, it may happen that the part of tree above current node has no colored nodes (all node values = 0) and in that case we cannot rely on the upper
                                                              // subtree to contrubute anytihing i.e. when we partition at current node we have find two sub-parts below the current node only such that one of them contain red and other contain atleast one black node
                                                             // So, in that case two things may happen :
                                                            // i) as we are going to remove the current node itself. so, all the sub-tree of the current nodes get separated for sure and so, if there exist two children tree such that one contains a red node and other a black
                                                           //     then we are done and in that case also we don't need to control anything i.e. all the possible subsets including current node is valid answer. This condition is denoted with the help of 'fg' variable (it checks
                                                          //      whether one of the previous children contains a black node and current children contains a red node OR one of the previous contains a red node and current children contains atleast one black node)
        
                                                        // ii) none of the two children exist such that one contains a red node and other one a black node. In this case we have to completely rely on one of the children such that the sub-sets it gives contains two parts such that one 
                                                       //      contains a red and other contains a black node. So, we consider only those sub-sets and now we already have our conditions meet and we are free to take any of the sub-sets from the other childrens.
                                                      //       Now, what we have to do is to iterate on each of the node and include correct ways from that sub-tree itself and take into account all the possible ways from the remaining children.
                                                     //        Doing it naively will take O(children^2) and there are a lot of cases (repitition to handle).
                                                    // So, we have to use Inclusion-Exclusion principle in that case also. Buttt.. we can do it in a smarter way!  What we do is to take all possible ways (all possible subset including current node) and just remove the ones that doesn't contain atleast a single children having True part!
        
                                                  // For, simplicity let's break the number of possible subset of a children tree into two parts (true + false) (true denotes the sets for which we don't need to rely on anyone else (i.e. this subtree only contains two partitions such that one contains a black and other a red node))
                                                 // Now, what we need to do is to take True part of a particular node and include all possible parts of other node and repeat the same for all the children one by one (this will contain repitition like "the case where two children gives their true part is counted twice" (try evaluating "(T2 + F2)*T1 + (T1 + F1)*T2" while the actual answer is only "T1F2 + T2F1 + T1T2")
                                                // so, computing this will may take (2^children) time if done naively! So, one more observation here : we have to exclude those part only in which there are no true part of any children (or other way : all possible ways of selecting only the false parts form all the children  AND this is exactly the same taks of computing all possible subsets)
    }
    
    if(color[s] == 1) b ++;
    else if(color[s]) r++;
    else{
        if(fg || (b && (tot_r - r)) || (r && (tot_b - b))) ans = (ans + sb) % mod, avail[s] = sb; // if upper tree or any of the two children(as a whole) have suficient nodes to rely on.
        else {
            ans = (ans + sb - x + mod) % mod;
            avail[s] = (sb - x + mod) % mod;
        }
    }
    if(color[s]) sb = 0;
}

void reset(){
    for(int j = 0; j <= n; j++){
        g[j].clear();
        red[j] = black[j] = sub[j] = avail[j] = 0;
    }
    ans = tot_r = tot_b = 0;
}

int main() {
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    
    int t;
    cin >> t;
    while(t--){
        reset();
        cin >> n;
        for(int j = 0; j < n-1; j++){
            int u, v;
            cin >> u >> v;
            g[u].pb(v);
            g[v].pb(u);
        }
        for(int j = 1; j <= n; j++){
            cin >> color[j];
            if(color[j] == 2) tot_r ++;
            else if(color[j]) tot_b ++;
        }
        dfs(1, 0);
        cout << ans << '\n';
    }
    return 0;
}

Show compiler warnings [ + ] Compiler args [ - ] Show input
Compilation time: 1.83 sec, absolute running time: 0.18 sec, cpu time: 0.11 sec, memory peak: 7 Mb, absolute service time: 2,02 sec