Hints: MEX Replacement on Tree

Before tackling the operation, can you efficiently compute $\sum f(v)$ on the original tree?

$f(v) = \mathrm{MEX}(S_v)$ where $S_v$ is the multiset of weights on the root-to-$v$ path. Since $p$ is a permutation, $S_v$ is a set of distinct values.

Think about what data structure you’d maintain during a DFS from the root, and how to query the MEX at each node.

Answer to Hint 1: Do a DFS from the root. When you enter node $v$, “mark” the value $p_v$ as present. When you leave $v$ (backtrack), “unmark” it.

At any point during the DFS, the marked values are exactly $S_v$ for the current node $v$. So you need a data structure that supports:

• Mark/unmark a value (point update)
• Query MEX (find the smallest unmarked value)

A segment tree over the value space $\{0, 1, \ldots, n\}$ does this. Store $1$ at position $x$ if $x$ is not on the current path, $0$ if it is. The MEX is the leftmost position holding $1$, which is the first prefix with nonzero sum. The max_right operation on a sum segment tree finds this in $O(\log n)$.

This gives you $\mathrm{mex}_1(v) = f(v)$ for every $v$ in $O(n \log n)$. Now think about what happens when you apply the operation.

The operation on vertex $v$: set $p_v \leftarrow f(v) = \mathrm{mex}_1(v)$.

Let $a = p_v$ (old weight) and $m = \mathrm{mex}_1(v)$ (the MEX of $v$’s root path). The operation replaces $a$ with $m$ at vertex $v$.

Which vertices $u$ have their $f(u)$ changed? Only those whose root-to-$u$ path passes through $v$, i.e., $u$ is in the subtree of $v$. For every other vertex, $p_v$ is not on its path, so nothing changes.

So the “gain” from operating on $v$ is: $\sum_{u \in \mathrm{sub}(v)} \bigl(\text{new } f(u) - \text{old } f(u)\bigr)$.

For any descendant $u$ of $v$, we have $S_v \subseteq S_u$ (the path to $u$ extends the path to $v$). This means every value present on $v$’s path is also present on $u$’s path.

What does this imply about the relationship between $\mathrm{mex}_1(v)$ and $\mathrm{mex}_1(u)$?

Answer to Hint 4: Since $\{0, \ldots, m-1\} \subseteq S_v \subseteq S_u$, the MEX of $S_u$ is at least $m$. Precisely:

• $\mathrm{mex}_1(u) = m$ if and only if $m \notin S_u$ (value $m$ never appears on the path from $v$ to $u$).
• $\mathrm{mex}_1(u) > m$ if and only if $m \in S_u$ (some node between $v$ and $u$ carries weight $m$).

This splits descendants of $v$ into two clean cases. Analyze what happens to $f(u)$ in each case when we swap out $a$ and swap in $m$.

If $\mathrm{mex}_1(u) > m$, then $m \in S_u$. After the operation, $S_u$ loses $a$ and gains $m$, but $m$ was already there. So effectively $S_u$ just loses $a$:

$$S_u' = S_u \setminus \{a\}$$

• If $a \ge \mathrm{mex}_1(u)$: the values $\{0, \ldots, \mathrm{mex}_1(u)-1\}$ are undisturbed, so $f(u)$ is unchanged. Change = 0.
• If $a < \mathrm{mex}_1(u)$: removing $a$ creates a gap at $a$. Change = $a - \mathrm{mex}_1(u)$ (always $\le 0$).

In this case, the contribution is never positive.

If $\mathrm{mex}_1(u) = m$, then $m \notin S_u$. The operation adds $m$ to $S_u$ and removes $a$:

$$S_u' = (S_u \setminus \{a\}) \cup \{m\}$$

Since $\{0, \ldots, m-1\} \subseteq S_u$ and $m \notin S_u$: adding $m$ fills the first gap.

• If $a < m$: removing $a$ creates a gap at $a < m$. New MEX $= a$. Change $= a - m < 0$.
• If $a > m$: the values $\{0, \ldots, m\}$ are now all present. The new MEX is somewhere $> m$. This is the only case with a positive change!

For the case $a > m$, what is the new MEX exactly?

Answer to Hint 7: When $a > m$, after adding $m$ and removing $a$, the set $S_u'$ has $\{0, \ldots, m\}$ present. The new MEX is the smallest gap starting from $m + 1$.

Define $\mathrm{mex}_2(u)$ as the smallest value $\ge m + 1$ that is not in $S_u$. (Equivalently, $\mathrm{mex}_2(u) = \mathrm{MEX}(S_u \cup \{m\})$.) You can compute this alongside $\mathrm{mex}_1$ during the same DFS.

Now the new MEX depends on whether removing $a$ hits before or after $\mathrm{mex}_2(u)$:

• If $a < \mathrm{mex}_2(u)$: the gap at $a$ comes first. New MEX $= a$. Change $= a - m$.
• If $a \ge \mathrm{mex}_2(u)$: the gap at $\mathrm{mex}_2(u)$ comes first. New MEX $= \mathrm{mex}_2(u)$. Change $= \mathrm{mex}_2(u) - m$.

In both sub-cases: Change $= \min(a, \mathrm{mex}_2(u)) - m > 0$.

Combining the cases, operating on $v$ (with $a = p_v$, $m = \mathrm{mex}_1(v)$) has:

• Positive contributions from descendants $u$ with $\mathrm{mex}_1(u) = m$ (Case B, $a > m$): each gives $+(\min(a, \mathrm{mex}_2(u)) - m)$.
• Negative contributions from descendants $u$ with $\mathrm{mex}_1(u) > a$ (Case A): each gives $+(a - \mathrm{mex}_1(u))$.
• Zero from the rest.

For $a < m$ (impossible since $m \notin S_v$ but $a = p_v \in S_v$; also $a = m$ is impossible for the same reason). So $a$ and $m$ are always distinct. But notice: $a$ could be less than $m$ if $v$’s weight contributes to a gap before $m$… no, actually $a \in S_v$ and $m = \mathrm{MEX}(S_v)$, so $a \ne m$ always. And $a$ can be either less than or greater than $m$.

But if $a < m$: all contributions are $\le 0$ (Case A gives $\le 0$; Case B gives $a - m < 0$). So we only want to operate on $v$ where $p_v > \mathrm{mex}_1(v)$.

For a candidate vertex $v$ with $a = p_v > m = \mathrm{mex}_1(v)$, the gain is:

$$\mathrm{gain}(v) = \underbrace{\sum_{\substack{u \in \mathrm{sub}(v) \\ \mathrm{mex}_1(u) = m}} \bigl(\min(a, \mathrm{mex}_2(u)) - m\bigr)}_{\text{positive part}} + \underbrace{\sum_{\substack{u \in \mathrm{sub}(v) \\ \mathrm{mex}_1(u) > a}} (a - \mathrm{mex}_1(u))}_{\text{negative part}}$$

The answer is $\sum_v \mathrm{mex}_1(v) + \max\bigl(0, \max_v \mathrm{gain}(v)\bigr)$.

These two sums can be computed separately. How would you compute the negative part efficiently for all $v$?

Answer to Hint 10 (negative part): Flatten the tree via Euler tour so each subtree corresponds to a contiguous range $[L_v, R_v)$.

For the negative part, you need: for each $v$ with $a = p_v$, the sum of $(a - \mathrm{mex}_1(u))$ over $u \in \mathrm{sub}(v)$ with $\mathrm{mex}_1(u) \ge a$.

Rewrite as: $a \cdot |\{u : \mathrm{mex}_1(u) \ge a\}| - \sum \mathrm{mex}_1(u)$ over those $u$.

Sweep values from $n - 1$ down to $0$. When you reach value $i$, insert all vertices $u$ with $\mathrm{mex}_1(u) = i$ into a segment tree (at their Euler-tour position). Then for the vertex $v$ with $p_v = i$, query its subtree range to get both the count and the sum of $\mathrm{mex}_1$ values. This handles the negative part in $O(n \log n)$.

For the positive part, you need: for each $v$ with $a = p_v > m = \mathrm{mex}_1(v)$, the sum of $\min(a, \mathrm{mex}_2(u)) - m$ over $u \in \mathrm{sub}(v)$ with $\mathrm{mex}_1(u) = m$.

A crucial observation: both $v$ and those descendants $u$ share the same $\mathrm{mex}_1$ value $m$. So group all vertices by their $\mathrm{mex}_1$ value and process each group independently.

Within the group for value $m$, each vertex can play two roles:

1. As an “operator” (= $v$): if $p_v > m$, it wants to query $\sum_{u \in \mathrm{sub}(v)} \min(p_v, \mathrm{mex}_2(u)) - m$.
2. As a “beneficiary” (= $u$): it contributes $\min(a, \mathrm{mex}_2(u)) - m$ to some ancestor $v$’s gain.

How do you handle the $\min(p_v, \mathrm{mex}_2(u))$ term efficiently?

Answer to Hint 12: Fix a group with $\mathrm{mex}_1 = m$. For “operator” $v$ with value $a = p_v$, the contribution of “beneficiary” $u$ (in $v$’s subtree) is:

$$\begin{cases} a - m & \text{if } \mathrm{mex}_2(u) > a \\\ \mathrm{mex}_2(u) - m & \text{if } \mathrm{mex}_2(u) \le a \end{cases}$$

Sort the group members by $p$-value (descending) for processing operators. Use a segment tree on the Euler tour. Initially, place $\mathrm{mex}_2(u) - m$ at each member $u$’s position.

As you process operators $v$ in decreasing order of $p_v$: “transition” all members $u$ with $\mathrm{mex}_2(u) > p_v$ — change their segment-tree entry from $(\mathrm{mex}_2(u) - m)$ to just $(-m)$, and also track a count. Then for operator $v$, the subtree sum plus $(\text{count}) \times p_v$ gives exactly $\sum \min(p_v, \mathrm{mex}_2(u)) - m$.

This is a standard “offline sweep + segment tree” pattern and runs in $O(n \log n)$ overall.

Algorithm summary:

1. DFS from root with a value-space segment tree. Compute $\mathrm{mex}_1(v)$, $\mathrm{mex}_2(v)$, and Euler tour $[L_v, R_v)$ for all $v$. Base answer $= \sum \mathrm{mex}_1(v)$.
2. Negative part (sweep from $n-1$ to $0$): insert vertices by $\mathrm{mex}_1$ value, query subtree sums for each $v$ at $p_v = i$. Accumulate into $\mathrm{gain}(v)$.
3. Positive part (group by $\mathrm{mex}_1$): within each group, sweep operators by decreasing $p$-value, transitioning beneficiaries by $\mathrm{mex}_2$. Accumulate into $\mathrm{gain}(v)$.
4. Answer $= \text{base} + \max(0, \max_v \mathrm{gain}(v))$.

Each phase is $O(n \log n)$, giving $O(n \log n)$ total.

• $\mathrm{mex}_1(v) = p_v$ is impossible ($p_v \in S_v$ but $\mathrm{mex}_1(v) \notin S_v$), so $a \ne m$ always.
• Vertex $v$ itself is in its own subtree: when computing the positive part, $v$’s own entry (with $\mathrm{mex}_1(v) = m$) participates as both operator and beneficiary — this is correct because the operation changes $f(v)$ too.
• The negative part captures vertices with $\mathrm{mex}_1(u) \ge a$ (including those equal to $a$ where the contribution is $0$), and the positive part captures vertices with $\mathrm{mex}_1(u) = m < a$. Together they cover all descendants.
• Clean up the segment tree after each mex group to avoid interference between groups.