PROBLEM LINK:
Author: Praveen Dhinwa
Tester: Misha Chorniy
Editorialist: Animesh Fatehpuria
PROBLEM
You 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:
- Each vertex in W is uncolored.
- W is a connected subgraph of T.
- If you remove all the vertices in W, there will be at least one pair of vertices (u, v) such that u is colored black and v red and there is no path from node u to node v.
Output your answer modulo 10^9 + 7. Constraint: N \le 10^5.
PREREQUISITES
Basic Familiarity with Tree DP.
EXPLANATION
There 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 different-colored nodes. We call a vertex v good iff v is uncolored and the removal of v disconnects at least one pair of different-colored 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 Nodes
Root 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:
- u is uncolored.
- There are at least two components with colored nodes in T - \{u\}.
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 Nodes
For 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.
Tester’s solution can be found here.