PROBLEM LINK:
Author: Praveen Dhinwa
Testers: Kevin Atienza, Sergey Kulik and Vasya Antoniuk
Translators: Vasya Antoniuk (Russian), Sergey Kulik (Russian), Team VNOI (Vietnamese) and Hu Zecong (Mandarin)
Editorialist: Kevin Atienza
DIFFICULTY:
Easy-medium
PREREQUISITES:
Combinatorics, dynamic programming
PROBLEM:
How many ways are there to color a tree of n nodes with k colors so that for every two nodes with the same color, all nodes in the path between them also have the same color?
QUICK EXPLANATION:
The lowest common ancestor of each color determines the coloring completely. In other words, once we select a subset of the nodes S and color them all differently, the color of every node a can be obtained by walking the path from a to node 1 and stopping at the first node in S. The color of a is the color of that final node.
Thus, the answer is the number of possible choices for S.
Note that the root must be the LCA of one of the colors. If there are i colors, then there are \binom{n-1}{i-1} ways to choose the non-root LCAs of the other colors, \binom{k}{i} ways to choose the colors, and i! ways to assign these colors to the LCAs. Thus, there are \binom{n-1}{i-1} \binom{k}{i} i! ways all in all. Summing across all valid i, we get:
Note that the answer is independent of the tree structure! It only depends on n and k.
To compute this, we precompute \binom{n}{r} and n! modulo 10^9 + 7 for all 0 \le n, r \le 50 using Pascal’s identity, and then compute the sum above (modulo 10^9 + 7) in O(k) time.
EXPLANATION:
Valid colorings
Observe that the constraint “for every two nodes with the same color, all nodes in the path between them also have the same color” is equivalent to the following: For every color, if you remove all nodes not colored with that color, along with the edges adjacent to them, then the remaining subgraph is connected. This means that every color induces a subtree!
We can visualize this in another way by rooting the tree at some node, say, node 1. In this case, if you view only the nodes that are of some color, then it also looks like a rooted tree, rooted at some node. This “root” is actually the lowest common ancestor (LCA) of all the nodes of that color.
What’s amazing is that, once we fix the lowest common ancestor for every color, the color of every other node is now completely determined! To see this, suppose we fix the nodes S = \{i_1, i_2, \ldots, i_m\} as the LCAs of the colors that appear in the tree. Then for every node i, its color must be the color of its lowest ancestor that is in S! This can be seen intuitively (e.g. to draw some examples to visualize this), but here’s a proof:
Suppose i_k is the lowest ancestor of i that is in S, and suppose \operatorname{color}(i_j) = \operatorname{color}(i) for some i_j \in S. Now, suppose for the purpose of contradiction that j \not= k. Then there are two cases, depending on whether i_j is a descendant of i_k or not:
- Suppose i_j is a descendant of i_k. Then i_j cannot be an ancestor of i, because i_k is the lowest ancestor of i in S. This means that LCA(i_j,i) \not= i_j, which contradicts the fact that i_j is the LCA of all nodes of the same color as i_j.
- Suppose i_j is not a descendant of i_k. Note that i is a descendant of i_k, which means that i_k is on the path from i to i_j, which fails our constraint.
In either case, we have a contradiction. Thus, j = k after all.
This means that the number of valid colorings depend on the number of choices for the set S of LCAs, and also on the colors that we give these nodes. But notice that we can choose the LCAs without actually looking at the edges! It means that *the answer is independent on the edges, and is only dependent on n and k!
Counting valid colorings
Now, we know what we need to count. We need to count the number of valid selections of the set S, and the number of assignments of colors among them. S can be any set of nodes, with one condition: node 1 must be in S. (Why?)
Now, let’s compute the answer. Suppose exactly i colors appear in the tree. How many choices for S are there? Well, S has i elements, but one element has already been chosen for us, namely node 1, so we need to select the remaining i-1 nodes among the remaining n-1 nodes. There are \binom{n-1}{i-1} choices. Next, we need to assign colors to these nodes. First, we must select i colors among the k available. There are \binom{k}{i} choices. Finally, how many ways are there to assign these colors to the nodes? There are i choices of colors for the first node, i-1 for the second node, i-2 for the third node, etc. Therefore, there are i(i-1)(i-2)\cdots 2\cdot 1 = i! assignments. Therefore, the number of colorings of a tree with exactly i colors is \binom{n-1}{i-1} \binom{k}{i} i!. Summing across all valid values of i, we get the following answer:
Computing this value modulo 10^9 + 7 can be done in O(k) time assuming we have already computed the binomial coefficients and factorials modulo 10^9 + 7. This can be done with dynamic programming:
The base cases are 0! = 1 and \binom{n}{0} = \binom{n}{n} = 1.
Since n and r can only reach up to 50 in this problem, we can quickly tabulate all the needed values.
Time Complexity:
O(n+k)