MEX Tree Manipulation
Contest: Codeforces Round 831 (Div. 1 + Div. 2)
Difficulty rating: 3300
Tags: data structures, trees

Given a rooted tree, define the value of vertex $$$u$$$ in the tree recursively as the MEX$$$^\dagger$$$ of the values of its children. Note that it is only the children, not all of its descendants. In particular, the value of a leaf is $$$0$$$.

Pak Chanek has a rooted tree that initially only contains a single vertex with index $$$1$$$, which is the root. Pak Chanek is going to do $$$q$$$ queries. In the $$$i$$$-th query, Pak Chanek is given an integer $$$x_i$$$. Pak Chanek needs to add a new vertex with index $$$i+1$$$ as the child of vertex $$$x_i$$$. After adding the new vertex, Pak Chanek needs to recalculate the values of all vertices and report the sum of the values of all vertices in the current tree.

$$$^\dagger$$$ The MEX (minimum excluded) of an array is the smallest non-negative integer that does not belong to the array. For example, the MEX of $$$[0,1,1,2,6,7]$$$ is $$$3$$$ and the MEX of $$$[6,9]$$$ is $$$0$$$.

Input format:
The first line contains a single integer $$$q$$$ ($$$1 \le q \le 3 \cdot 10^5$$$) — the number of operations.

Each of the next $$$q$$$ lines contains a single integer $$$x_i$$$ ($$$1 \leq x_i \leq i$$$) — the description of the $$$i$$$-th query.

Output format:
For each query, print a line containing an integer representing the sum of the new values of all vertices in the tree after adding the vertex.

Constraints:
Time limit: 5.0 s
Memory limit: 512.0 MB

Note:
In the first example, the tree after the $$$6$$$-th query will look like this.

- Vertex $$$7$$$ is a leaf, so its value is $$$0$$$.
- Vertex $$$6$$$ is a leaf, so its value is $$$0$$$.
- Vertex $$$5$$$ only has a child with value $$$0$$$, so its value is $$$1$$$.
- Vertex $$$4$$$ is a leaf, so its value is $$$0$$$.
- Vertex $$$3$$$ only has a child with value $$$0$$$, so its value is $$$1$$$.
- Vertex $$$2$$$ has children with values $$$0$$$ and $$$1$$$, so its value is $$$2$$$.
- Vertex $$$1$$$ has children with values $$$1$$$ and $$$2$$$, so its value is $$$0$$$.

The sum of the values of all vertices is $$$0+0+1+0+1+2+0=4$$$.

Example input:
7
1
1
3
2
5
2
1

Example output:
1
1
3
2
4
4
7

Example input:
8
1
1
1
1
5
6
7
8

Example output:
1
1
1
1
3
2
4
3

Editorial:
Define $$$\text{MEX}_i(A)$$$ of a certain array $$$A$$$ as the $$$i$$$-th smallest non-negative integer that does not belong to the array.
Let's solve a simpler problem. Imagine that you are given an array $$$A$$$, you are asked to append a new element $$$x$$$ to the array, and you need to calculate the MEX of the array after the operation. It is easy to see that the new value of MEX is either $$$\text{MEX}_1(A)$$$ if $$$x \neq \text{MEX}_1(A)$$$ or $$$\text{MEX}_2(A)$$$ if $$$x = \text{MEX}_1(A)$$$ where $$$A$$$ is the array before the operation.
Let $$$(x, y, z)$$$ define a relationship between input and output which means:
• If the input is not $$$x$$$, then the output is $$$y$$$.
• If the input is $$$x$$$, then the output is $$$z$$$.
This means the relationship between the new element to be appended to $$$A$$$ and the new value of MEX can be expressed as $$$(\text{MEX}_1(A), \text{MEX}_1(A), \text{MEX}_2(A))$$$. Define that triple as $$$\text{MEXtuple}(A)$$$.
Define the heavy child of vertex $$$u$$$ as the child with the largest subtree size. If there are children with the same subtree size, choose one of them. Let's call each child other than the heavy child as a light child. In particular, the edge that connects a vertex to its heavy child is called a heavy edge.
Let $$$C_u$$$ be the array that contains the values of the light children of vertex $$$u$$$. We can see that if we maintain $$$\text{MEXtuple}(C_u)$$$ we can get the relationship between the value of the heavy child and the value of vertex $$$u$$$.
Let's consider a path in the tree that traverses through parents while only using heavy edges. Consider the relationship between the value of the vertex at the beginning of the path and the end of the path. Let's calculate the $$$\text{MEXtuple}$$$ for each vertex in the path other than the first one. We can see that we can get the relationship between the vertex in the beginning of the path and the end of the path using those $$$\text{MEXtuple}$$$ values. In fact, that relationship can also be expressed as a triple $$$(x, y, z)$$$ defined above.
Let's decompose the tree into chains of heavy edges. For each chain, construct a segment tree with each node in the segment tree calculating the relationship triple $$$(x, y, z)$$$ for a segment of the chain.
Suppose we want to add a new vertex. The vertices that have their values changed are the ones that lie on the path from the new vertex to the root. Observe that that path only traverses through at most $$$O(\log Q)$$$ different heavy edge chains. For each chain, we can get the value of the vertex at the end of the chain by doing a point update on the chain's segment tree. When we move between two different chains, we update the $$$C_u$$$ for a vertex and recalculate it's $$$\text{MEXtuple}$$$ in $$$O(\log Q)$$$ time complexity using a set or another segment tree. We repeat this process until we reach the root. We can calculate that the total complexity of a single query is $$$O(\log^2 Q)$$$.
To get the sum of the values of the entire tree after an operation, we can store two new variables in each node of the segment tree of each chain that represents the sum of the chain segment for the two possible cases of the input.
Time complexity: $$$O(Q \log^2 Q)$$$