PROBLEM LINK :
Setter : Vishal Mahavar
Testers : Prasann Kumar Gupta, Onkar Ratnaparkhi
DIFFICULTY :
Medium-Hard
PREREQUISITES :
Trees, Centroid decomposition on trees
PROBLEM
Given a tree with N nodes and N-1 edges, where each node can have two colors represented by 0 or 1. We have to process Q queries of two types on it :
- 1\> x - Change the color of node x (from 0 to 1 and vice versa).
- 2\>x\>k\>c - Find the number of nodes with color c at a distance k from node x.
EXPLANATION
Since the queries are only related to the distance between the nodes and not the actual path, it’s intuitive to come up with a centroid decomposition solution.
We are going to use is the following properties of the centroid tree :
- Height of the centroid tree is O(logN).
- Let sz_i denote the size of the subtree of centroid i in the centroid tree, then \sum\limits_{i=1}^N sz_i is O(NlogN).
- For any pair of nodes p and q, the path from p to q in the original tree can always be broken into the concatenation of paths from p to L and L to q. Here L is the lowest common ancestor of nodes p and q in the centroid tree.
Let’s denote dist[i][c][d] as number of nodes which have color c and are at a distance of d from centroid i in it’s subtree. It’s easy to see that the number of distinct values of d can be at most sz_i, hence it’s feasible to compute it (second property).
Also consider the following computation, dist\_par[i][c][d] which is the same as the above computation but the difference is that we are measuring distance from the parent of i in the centroid tree instead of the centroid i. Obviously this is not needed to be computed for the root of the centroid tree. These computations can be done exactly the same way we process updates. Let’s see how we can process updates and queries using the above two computations.
Update
WLOG, consider a node x whose color is 0 and is to be changed to 1. The values of dist[y][][] will change only for the ancestors of this node in the centroid tree (including the node itself). So, we traverse up the centroid tree (first property) and change this for every centroid y. Let’s say distance between p and q nodes (in the original tree) is denoted using d_{xy}, then dist[y][0][d_{xy}] will reduce by one and dist[y][1][d_{xy}] will increment by one.
Similarly if y is not the root of the centroid tree, let z be the parent of y. Then dist\_par[z][0][d_{zx}] will reduce by one and dist\_par[z][1][d_{zx}] will increment by one.
Query
Once again, let’s assume we need number of nodes at distance of k from node x which are of color 0. Using the third property mentioned above, we only need to go through the ancestors of the node x in centroid tree. For each such ancestor y, we need to find the number of required nodes which are at distance k - d_{xy} from y. Which can be written as dist[y][0][k-d_{xy}] and be added to our answer.
But here’s a catch. Consider the node z such that it is a child of y in the centroid tree and is also an ancestor of node x. The nodes present in the subtree of z cannot be included using the third property (considering node y as the breakpoint of paths). The reason is, the LCA of such nodes and the node x is node z, NOT node y hence these nodes have to be counted when we are at node z. In order to remove such nodes, we need to subtract dist\_par[z][0][k-d_{xy}] from our answer.
We are traversing O(logN) on the centroid tree. Since we only need to store the required d's in dist[i][c][d], we’ll use a map (in C++) for this purpose which gives an extra logN factor to the traversal. And we can use both logarithmic time and constant time method for calculating distances between any two nodes in the tree, they will be added paralelly to the time complexity. Hence the time complexity per query is O(log^2N). The preprocessing itself is same as processing updates on N nodes. The total time complexity is O((N+Q)log^2N).
NOTE: You can also use unordered_map(or any other hashing based data structure) to solve this in O((N+Q)logN).
SOLUTIONS
Setter’s Solution
Tester’s Solution
Feedbacks and suggestions are always welcomed.
Do discuss the solution/any other approach in comments 
.