PROBLEM LINK:Author: Praveen Dhinwa PROBLEMYou are given a tree $T$ with $N$ vertices. A vertex can be either black, red, or uncolored. There is at least one black and at least one red vertex. Compute the number of subsets of vertices $W$ such that:
Output your answer modulo $10^9 + 7$. Constraint: $N \le 10^5$. PREREQUISITESBasic Familiarity with Tree DP. EXPLANATIONThere are many possible solutions to this problem, all of which use some sort of dynamic programming on trees. Thus, it is required to have some level of familiarity with that topic. If you are not familiar with DP on Trees, check out this blog. We want to compute number of connected subgraphs of uncolored vertices that disconnect at least one pair of differentcolored nodes. We call a vertex $v$ good iff $v$ is uncolored and the removal of $v$ disconnects at least one pair of differentcolored nodes. We want to compute number of connected components in $T$ of uncolored vertices that contain at least one good node. This is equivalent to computing all possible connected subsets, and then subtracting those connected subsets (of uncolored nodes) that have no good nodes. It turns out that this can be done with a straightforward dynamic program. Precomputing all Good NodesRoot the tree arbitrarily. Compute for each node $v$ two quantites, $red[v]$ and $black[v]$ denoting the number of red and black vertices in the subtree rooted at $v$ respectively. This can be done with one DFS. To check if a node $u$ is good, consider $T  \{u\}$. A node $u$ is good iff:
Thus, we can check all its neighbours of $u$ and see if there are at least two neighbours whose components have colored vertices by using our precomputed lists $red[v]$ and $black[v]$. Using this approach, we can precompute all good nodes in amortized $O(n)$ time. Computing Number of Connected Subgraphs with no Good NodesFor this step, we require to do a tree DP. Again, we root the tree arbitrarily, say at node $r$. Let $dp[v][0]$ denote the number of valid subgraphs in the subtree rooted at $v$ that don't include $v$, and let $dp[v][1]$ denote the number of valid subgraphs in the subtree rooted at $v$ that include $v$. The answer is $dp[r][0] + dp[r][1]$. The transitions are fairly straightforward and I'll leave it as an exercise. As a hint, note that the transition depends on the color of the node. If you are struggling to find these transitions, please go through the blog mentioned above. It explains the basic ideas of Tree DP with some very good examples. If you're still stuck, please read some available solution. Since there are $O(n)$ states (and the transition is $O(1)$), the total time complexity is $O(n)$. AUTHOR'S AND TESTER'S SOLUTIONS:Author's solution can be found here.
This question is marked "community wiki".
asked 28 Dec '17, 14:37
