How many colourings
The debugger of a red–black tree library saves snapshots of a tree without colours, only its shape: which children each node has. To know whether the colouring can be recovered from such a snapshot, we need to know how many colourings of the shape obey the rules: zero means no red–black tree has this shape, one means the colours are recovered uniquely, more means the snapshot is ambiguous.
A colouring is a choice of colour, red or black, for every node. It is valid if the root is black, a red node has no red children, and all paths from the root down to the empty links have the same number of black nodes. The keys in the nodes do not matter, only the shape. Tests three, four, five and nine are shapes of real trees: an AVL tree of the 19,073 earthquakes of the USGS catalogue inserted in time order, a red–black tree and a plain search tree of the first 20,000 distinct words of War and Peace, and an AVL tree of legost.in page addresses. The node numbers are shuffled.
Input
The first line holds n and r (1 ≤ n ≤ 10⁵, 1 ≤ r ≤ n): the number of nodes and the number of the root. Then n lines: the i-th holds the numbers of the left and right child of node i, 0 if there is no child. It is guaranteed that this is a binary tree with root r containing all n nodes.
Output
The number of valid colourings modulo 10⁹ + 7.
Examples
Example 1
3 1 2 3 0 0 0 0
2
Example 2
3 1 2 0 3 0 0 0
0
Limits
- Time · Python
- 4000 ms
- Time · Go
- 1000 ms
- Memory
- 256 MB
- Judging
- 11 tests
Walkthrough
Open it step by step: each step gives away a little more. You can jump straight to the solution, but a problem solved with one hint stays with you longer.
1Hint 1
Trying all colourings is 2n options, but the rules are checked bottom-up: about a subtree it is enough to know the colour of its root and how many black nodes lie on any path down from the root. Try counting the colourings of each subtree for each such number.
2Hint 2
Let the subtree of node v have black height b: that many black nodes, counting v, on any path down from v. If v is black, both children, of either colour, have black height b − 1. If it is red, both children are black with black height b. An empty link is black with black height 0.
3Hint 3
Trying every b from 0 to n for each node takes too long. Which values of b are possible at all if the shortest path from v down to an empty link contains s nodes and the longest ℓ?
4The idea
Dynamic programming on the tree, bottom-up. For each node v and black height b keep two numbers: B(v, b), the colourings of the subtree with v black, and R(v, b), with v red. Then B(v, b) = (B + R)(left, b − 1) · (B + R)(right, b − 1), R(v, b) = B(left, b) · B(right, b), and an empty link has B = 1 at b = 0. The answer is the sum of B(root, b) over all b. The possible b lie between ⌊ℓ/2⌋ and s, and in a balanced tree s is at most log₂(n + 1).
5Walkthrough
Why the formulas are right. The rules of a red–black tree are checked subtree by subtree: a colouring of the subtree of v is valid (apart from the condition on the root) if the colourings of the children’s subtrees are valid, a red v has no red children, and the children’s black heights are equal. The black height of v is the children’s black height plus one if v is black. The colourings of the left and right subtrees are chosen independently, so the counts multiply, while different colours of a child add up. The condition “the root is black” is applied at the very end: only B is summed.
Which b are possible. The shortest path down from v has s nodes, and all the black nodes of the path are among them, so b ≤ s. The longest path has ℓ nodes, reds do not stand in a row, and each red is followed by a black (the empty link at the end is black too), so there are at most ⌈ℓ/2⌉ reds and at least ⌊ℓ/2⌋ blacks among the path’s nodes. For each node keep an array only for b from ⌊ℓ/2⌋ to s; if this range is empty, there are no colourings. When ℓ exceeds 2s + 1 at even one node, the answer is zero: this happens with the plain search tree of the novel’s words, whose root has a shortest path down of 5 nodes and a longest of 33.
The traversal order. Children must be computed before the parent. Recursion on a path of 10⁵ nodes (the tenth test) would overflow Python’s stack, so the order is built with an explicit stack: a top-down traversal, then a pass over it in reverse.
What comes out on real shapes. AVL trees can be coloured in a great many ways: modulo 10⁹ + 7 the answer looks like a random number. The red–black tree of the novel’s words can of course be coloured: the insertion itself gave it one colouring. And the skinniest AVL tree of height 23, the Fibonacci tree of 75,024 nodes, can be coloured in exactly one way: at each of its nodes the longest path down is almost twice the shortest, and there is no choice of black height.
5+Solution
6Complexity
A node stores at most s − ⌊ℓ/2⌋ + 1 values, each computed in O(1). In a tree of n nodes s at the root is at most log₂(n + 1), otherwise the top s levels would not have enough nodes, and the same holds at any node. In total O(n log n) time and memory, and in practice much less: most nodes have small subtrees.
7Common mistakes
- A red under a red. If a red node is allowed children of any colour, the answer comes out larger than the correct one, already on the third test.
- A red root. By the statement the root is black; adding B and R at the root also inflates the answer.
- Recursion. A recursive traversal in Python will not survive a path of 10⁵ nodes.
- All b in a row. An array of length n per node is 10¹⁰ operations; only the b between ⌊ℓ/2⌋ and s are needed.
Algorithms behind it
O(log n) for search, insertion and deletion
2 problems
DP
Technique
Dynamic programming on trees
A subtree’s answer from its children’s answers: a pass from the leaves to the root, and for all roots at once, rerooting.
O(n)
2 problems