Problem Link: contest, practice
Difficulty: Easy-Medium
Pre-requisites: DP on Trees
Problem:
Given a rooted tree with N vertices. You would like to paint each vertex of the tree in one of three possible colors: red, green, and blue. The following conditions must be fulfilled: for each vertex painted red, there must be no more than R red vertices in its subtree; for each vertex painted green, there must be no more than G green vertices in its subtree; for each vertex painted blue, there must be no more than B blue vertices in its subtree.
Find the number of ways to paint the tree modulo 1000000007.
Explanation:
First of all, I strongly recommend you to read the following article about calculating knapsack problem on a tree(link; check out Chapter 7.3). Otherwise, it’s gonna be super-hard to understand this editorial.
This is original Author’s explanation(corrected and edited by me). It’s quite tricky and there’s no implementation details, but the main idea is described. If you need details - feel free to ask them in the comments.
After that, if we have some vertex V painted in some color, and this vertex is OK, than we don’t care about vertices of the same color in the subtree of V.
Then we should handle two DP functions. Let’s assume, that the root is painted in red color(it can be blue or green as well, but we’ll consider the only case; the others are similar).
The first DP looks like this: F(V)(cntRed). In this DP we ensure, that the number of red vertices in the subtree of V is OK. So, we’ll iterate though all the children of vertex V and will choose whether to paint them in red or in some other color. Let’s fix some child U. If we choose red color, then we call F, otherwise, we call G( G(V)(cntGreen) if we paint the child blue, and G(V)(cntBlue) if we paint the child green(not vice-versa) ).
G looks like this: G(V)(cntGreen) or G(V)(cntBlue). I will describe the first one. In this DP, we ensure that all green vertices are OK. Also note that when we call this DP from F, we know that the root(of our subtree) is painted blue. This means that we don’t need to care about all blue vertices in the subtree, only about the root (so only the number of blue vertices is interesting for us). But we have to keep the number of green vertices as the state of the DP. So we actually forget about the red color in G.
Now, when we made a call of G from F, we have some cntGreen vertices painted green, and Size(v) - cntGreen vertices that we painted blue (we actually brute force cntGreen). But, we can paint some of the blue vertices in red color. So, we also brute this number, and use some Cnk to find the number of ways to repaint some blue vertices in red color. Note, that we have to ensure that the number of blue vertices in the subtree of the child is also OK.
Total complexity is O( N3 ), O( N2 ) for calculing G and O( N3 ) for calculing F.
Please, check out Setter’s and Tester’s solutions for your better understanding.
Setter’s Solution: link
Tester’s Solution: link