PROBLEM LINK:
Author: Ramazan Rakhmatullin
Testers: Lewin Gan, Marek Sommer
Editorialist: Pawel Kacprzak
DIFFICULTY:
Hard
PREREQUISITES:
DFS, DFS order, trees, segment tree or treap
PROBLEM:
Let’s reformulate the problem a little bit first. For a given tree T with N nodes numbered from 0 to N-1, where each node has initially assigned color c_i, the goal is to handle Q queries numbered from 0 to Q-1. The i-th query contains a single integer x_i changes the color of node with number i \mod N to (x_i + prevAns) \mod N, where prevAns is the answer to query i-1. The answer for a query is defined as follows:
For each node v, let A[v] be the set of all ancestors of v with the same color as v in the tree. Similarly, let S[v] be the set of all successors of v with the same color as v in the tree. Then, f[v] is defined as A[v] \cdot S[v]. It counts exactly the number of triples of nodes (a, v, b), such that they all have the same color and a is an ancestor of v and v is an ancestor of b. Then, the answer for the i-th query is \sum\limits_{v \in T} f[v] after the color changes from the first i queries are made. The goal is to print answer before all queries are made and then to print the answer after each made query.
EXPLANATION:
Subtask 1
In the first subtask, both N and Q are at most 2000, which allows solving the problem in O(N) time for a single query, independently of the other queries.
The idea is to use DFS and smart counting. Let’s define two arrays:
g[c] := the current number of nodes of color c in the path to the root
dp[c] := the number of nodes of color c discovered so far
When we enter node v in DFS, we want to explicitly count the value of f[v] and add it to the total result. First, let’s save the current value of g[c_v] to a variable x and the current value of dp[c_v] to variable y. Next, we increment both g[c_v] and dp[c_v] taking node v into account and then we visit recursively all children of v using the same DFS. After that, values of g[c_v] and dp[c_v] are updated by the recursive calls, and the value of f[v] can be computed as x \cdot (dp[c_v] - y - 1). When we exit v's subtree, we decrement the value of g[c_v].
It immediately follows that this approach has overall O(N \cdot Q) time complexity.
Subtask 2
This subtask is mainly intended for methods with time complexity of O(Q \cdot \sqrt N), for example, some sqrt-decomposition approaches might be applied instead of segment tree used in the solution described for the last subtask.
Subtask 3
Basically, this subtask is designed for not efficient enough implementation of the method described for the last subtasks. However, there also exists methods of O(N \cdot \log N \cdot \log N) complexity, for example, based on persistent segment trees.
Subtask 4
In the last subtask, we have the original constraints, i.e N, Q \leq 5 \cdot 10^5. A really fast solution is needed to pass this subtask. There are two key observations that can be made:
- The tree is structured in such a way that the parent of node v has always a smaller number than v.
- If we divide the queries into the blocks containing exactly N queries, then consecutive queries in a single block change colors of nodes starting from node 0 and ending in node N.
Fact 1:
The above two observations imply that
in a single block of queries, for any
node v, all updates to ancestors of
v are always made before the update
to v.
Let’s assume that we are processing a fixed block of queries and the current query changes the color of node v and let’s see how this change of color affect the current result for the whole tree.
First of all, the process of changing the color of v from color_{old} to color_{new} can be divided into first subtracting the number of triples of nodes with color color_{old} containing v from the current result and then adding the number of triples of nodes with color color_{new} containing v after the color of v is updated to the updated result. In the editorial, we’re going to describe the process of subtracting the number of triples of nodes with color color_{old} containing v from the current result. The process of adding triples with color color_{new} containing v to the result can be performed almost analogously.
First of all, v appears exactly f[v] number of times as the middle node of a triple, so if we keep the value of f[v] up to date after each query, then we can first just subtract this value from the current result. Then, we are going to count the number of triples in which v appears as the first and as the last node.
The number of triples containing v as the first node can be computed as:
Let’s call these triple, the triples of the first kind.
Similarly, the number of triples containing v as the last node can be computed as:
Let’s call these triple, the triples of the second kind.
The key observation is that Fact 1 defined above allows us to efficiently compute the number of triples of both kinds.
Computing the number of triples of the second kind.
Let p[v] be the number of nodes of the color color_{old} at the path from the root of the tree to v before processing the current block of queries.
Then, let’s redefine S[v] as the set of nodes u such that v is an ancestor of u and u has color color_{old} before processing the current block of queries. Then, the number of triples of the second kind can be computed as:
Notice that the 1 is subtracted in order not to count node v itself. So, the problem of counting these triples is reduced to the problem of counting 3 values: \sum\limits_{u \in S[v]} p[u], |S[v]|, and p[v].
The value of p[v] can be updated while performing DFS down and up the tree, and both the sum of p[u] over S[v] and the size of S[v] can be computed by keeping for each color a vector of appropriate values for every node of that color along with its index in DFS ordering. This index can be used to perform a binary search over this vector to get what we want. For example, let’s assume that we want to compute the sum of p[u] over S[v]. Then in the vector for color color_{old} we store the pairs $(\text{prefix sum of p[x] for nodes x visited not later than v}, dfsVisitTime[v])$. Now, in order to get the desired sum, we can just perform a binary search over this vector to find the appropriate indices and calculate the desired value. Notice that the vector is, of course, ordered automatically by the visit times during DFS. In the same manner, we can compute the cardinality of S[v].
Computing the number of triples of the first kind.
Let q[v] be the number of nodes of the color color_{old} at the path from the root of the tree to v after processing the current block of queries. This means that q[v] contains all already updated nodes in this block on the path from the root to v.
Again, let’s redefine A[v] as the set of nodes u such that u is ancestor of v and u has color color_{old} after processing the current block of queries, or alternatively, after processing all queries in the current block before the current query. Then, the number of triples of the first kind can be computed as:
The value of q[v] can be updated similarly as p[v] while performing DFS over the tree. The cardinality of A[v] and the sum of q[u] over A[v] has to be now computed as values on the path from v up to the root rather than in the v's subtree as we did for counting the triples of the second kind. The following very useful trick can be used whenever we want to calculate a partial sum of some values on a path from v to the root:
We’re going to store an array and a segment tree over this array for range queries and point updates. When we enter node v with DFS, we add to first unused position in this array the desired value for node v, and when we exit node v, we add to the first unused position in this array the same desired value for v but with a negative sign. When we need to find a sum of values in a path from v to the root, we just take the partial sum from the beginning of the array to the time when we entered v.
In our problem, we keep such segment tree for each color in the tree and use values of q[u] as the elements of the underlying arrays. Similarly, we can compute the cardinality of the set A[v].
As stated previously, the above approach subtracts all disappearing triples from the previous result after the color of v is changed. An analogous method can be used to count all the triples that are created when the color_{new} is assigned to v. Both these operations keep the final result up to date after each update. The total time complexity of this method is O((N+Q) \cdot \log(N)), because for each single query, we do O(\log(N)) operations. For the implementation details of this exact method please refer to setter’s solution linked below. The only difference is that he used a treap in places when the editorial uses a segment tree.
AUTHOR’S AND TESTER’S SOLUTIONS:
Setter’s solution can be found here.
First tester’s solution can be found here.
Second tester’s solution can be found here.