You are not logged in. Please login at to post your questions!


COLTREE - Editorial




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




Combinatorics, dynamic programming


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?


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:
$$\sum_{i=1}^k \binom{n-1}{i-1} \binom{k}{i} i!$$ 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.


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:
$$\sum_{i=1}^{\min(n,k)} \binom{n-1}{i-1} \binom{k}{i} i!$$

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:
$$\begin{align*} n! &= n\cdot (n-1)! \\\ \binom{n}{r} &= \binom{n-1}{r-1} + \binom{n-1}{r} \end{align*}$$ 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:




This question is marked "community wiki".

asked 19 Jun '16, 15:16

kevinsogo's gravatar image

accept rate: 11%

edited 21 Jun '16, 23:06

admin's gravatar image

0★admin ♦♦

Great editorial,thank you!


answered 13 Apr '17, 02:30

ps_1234's gravatar image

accept rate: 0%

@admin Problem is not visible in Practice section

U.P.D : Its visible now Thank you


answered 21 Jun '16, 20:07

tihorsharma123's gravatar image

accept rate: 15%

edited 21 Jun '16, 20:58

toggle preview

Follow this question

By Email:

Once you sign in you will be able to subscribe for any updates here



Answers and Comments

Markdown Basics

  • *italic* or _italic_
  • **bold** or __bold__
  • link:[text]( "title")
  • image?![alt text](/path/img.jpg "title")
  • numbered list: 1. Foo 2. Bar
  • to add a line break simply add two spaces to where you would like the new line to be.
  • basic HTML tags are also supported
  • mathemetical formulas in Latex between $ symbol

Question tags:


question asked: 19 Jun '16, 15:16

question was seen: 3,724 times

last updated: 13 Apr '17, 02:30