Author: Sunny Aggarwal
Tester: Sergey Kulik
Editorialist: Florin Chirica
You’re given an infinite complete binary tree. If a node has a color, its sons have complementary color. You need to perform some operations: either swap colors or answer how many nodes have a given color on a path.
The idea is the distance from a node x to the root is at most log2(x). Also, we can notice the tree is uniquely determined by the color of the root. Using those observations, we can keep a variable for the color of the root and apply brute force for solve queries.
It’s all about parities
Let’s suppose the root of the tree has a color. Then, its sons will have different color. Also, the son of sons will have different color compared to sons. Since there are only two colors, this means sons of sons will have the same color as root.
Let’s define level of a node by distance from the root to that node. The only node at level 0 is root. The nodes at level 1 are the sons of root. The nodes from level one have different color from the root. The nodes from the level 2 (sons of sons) have the same color of the root. Also, the nodes from level 3 have different color from the root.
We can observe that if the parity of the level is 0, node has same color as the root. Otherwise, node has different color by the root.
Number of nodes from a path having a given color
This is a standard problem: count number of nodes from a given path having some property. Usually, in a tree to count something from a path from x to y, you count from the root to x, then add result from the root to y and subtract result from the root to LCA(x, y).
count(x, y, color) = count(root, x, color) + count(root, y, color) - 2 * count(root, LCA(x, y), color) + (1 if color(LCA(x, y)) is equal to color).
By LCA(x, y) I noted lowest common ancestor between x and y. Why does it work the formula? It’s clearly that count(root, x, color) + count(root, y color) counts the nodes from the path from x to y. However, it counts something more… it counts also number of nodes from root to LCA(x, y). It counts this once calling count(root, x, color) and once calling count(root, y, color). So, we need to subtract 2 * count(LCA(x, y), color). Now, it is still not good… Node LCA(x, y) is not counted at all now! However, it should be counted exactly once. So we can check if node LCA(x, y) has wanted color and if so, increase the result of count by 1. Now, the formula should be good.
Subproblem: how to implement function count
Let’s suppose we get all nodes from path from root to node x together with their levels. Then, it’s simple to count how many nodes have a given color.
- If root has the same color as color variable, we need to count how many nodes have parity of level even.
- Otherwise, we need to count how many nodes have parity of level odd.
All we need to do is, for a node x, to find its parent. Then, repeat the procedure for its parent. We know that node x is either 2 * k or 2 * k + 1. So k = x / 2. After 1 step, parent becomes x / 2. After two steps, parent becomes (x / 2) / 2 = x / 4. After s steps, x becomes x / 2 ^ s. Since our root is node 1, we’re interested in all nodes such as x / 2 ^ s >= 1. This means x >= 2 ^ s. Since x <= 10^9, it means 2 ^ s <= 10^9. So we get worst case, s = log2(10^9). This means, going up in the tree takes at most log2(10^9) ~ 31 steps until it reaches the root.
Since we do at most 31 steps, we can afford to find all nodes from path from node x to the root. We can also afford to find level of node x. Then, after each step, level will decrease by one, so it will be found in O(1). Since we have nodes and levels, we can easily count how many nodes have a given color from path from root to node x.
Handle operations Qi
We can keep a variable rootColor, initially BLACK. When you apply a “Qi” operation, it’s like flipping color of the rootColor (BLACK -> RED and RED -> BLACK). All tree is uniquely determined by the color of the root, so it’s enough to simply flip rootColor variable.
Subproblem: find LCA(x, y)
This part is our last challenge to finish the task. We calculate LCA like doing a brute force. Please refer to official solution to see the details.
Operation “Qi” takes only O(1) time, since you simply flip a variable. Operations “Qr” and “Qb” take O(log(max(x, y))) time, needed for going up in the tree.
AUTHOR’S AND TESTER’S SOLUTIONS:
To be updated soon
To be updated soon