### 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.