Run Code

API

Code Wall

Users

Misc

Feedback

Login

Theme

Privacy

Patreon
BLREDSET
#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 subtree 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 subtree rooted at i. avail[i] == number of valid ways to choose a set of uncolored nodes from the subtree 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 subtree 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 subtree 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 subparts 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 subtree 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 subsets it gives contains two parts such that one // contains a red and other contains a black node. So, we consider only those subsets and now we already have our conditions meet and we are free to take any of the subsets from the other childrens. // Now, what we have to do is to iterate on each of the node and include correct ways from that subtree 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 InclusionExclusion 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 < n1; 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; }
run

edit

history

help
0
Please
log in
to post a comment.
str_to_int
suffle array
Insertion Sort
Making pyramid using nested loop 2/2
nobita's candies problem
ammmma
Graph DFS
IviewwithSelf
const test
Policy based smart pointer
Please log in to post a comment.