# 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 .