Editorial : MEX Replacement on Tree

We are given a rooted tree with a permutation $p$ of $\{0, 1, \ldots, n-1\}$ as vertex weights. For each vertex $v$, define $f(v) = \mathrm{MEX}(S_v)$ where $S_v$ is the set of weights on the path from root to $v$. We may perform at most one operation: pick a vertex $v$ and set $p_v \leftarrow f(v)$. Maximize $\sum f(v)$.

The path from root $1$ to vertex $v$ passes through a sequence of vertices. $S_v$ is the set of their $p$-values, and $f(v) = \mathrm{MEX}(S_v)$, the smallest non-negative integer not in that set.

Root the tree at vertex $1$ and DFS. Maintain a segment tree over the value space $\{0, 1, \ldots, n\}$. Position $x$ stores $1$ if value $x$ is absent from the current root-to-node path, and $0$ if it is present.

• Enter node $v$: set position $p_v$ to $0$ (mark it present).
• Compute $f(v)$: query max_right(0) with predicate “sum $= 0$”. This finds the first position whose prefix sum is nonzero, i.e., the first absent value. That is $\mathrm{MEX}(S_v)$.
• Leave node $v$: set position $p_v$ back to $1$ (unmark it).

Since the DFS enters and leaves each node exactly once, this computes all $f(v)$ in $O(n \log n)$.

Base answer $= \sum_{v=1}^{n} f(v)$ when no operation is used.

When we operate 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). Note that $a \ne m$ always, because $a \in S_v$ but $m \notin S_v$.

Only vertices $u$ whose root-path includes $v$ — i.e., vertices in the subtree of $v$ — are affected. For such $u$, the weight $a$ on their path is replaced by $m$:

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

The final answer is $\text{base} + \max\bigl(0,\; \max_v \mathrm{gain}(v)\bigr)$.

Flatten the tree with a DFS Euler tour: each subtree of $v$ maps to a contiguous range $[L_v, R_v)$. This lets us convert subtree queries into range queries on a segment tree.

Since $S_v \subseteq S_u$ for any descendant $u$ of $v$, the values $\{0, \ldots, m-1\}$ are all in $S_u$. So $\mathrm{mex}_1(u) \ge m$. More precisely:

• $\mathrm{mex}_1(u) = m$ $\iff$ $m \notin S_u$ (nobody on the $v$-to-$u$ path carries weight $m$).
• $\mathrm{mex}_1(u) > m$ $\iff$ $m \in S_u$.

This gives a clean split into two cases.

Adding $m$ to $S_u$ does nothing (it was already there), so effectively we just remove $a$:

• If $a \ge \mathrm{mex}_1(u)$: the gap structure below $\mathrm{mex}_1(u)$ is untouched. Change $= 0$.
• If $a < \mathrm{mex}_1(u)$: removing $a$ creates a new gap. Change $= a - \mathrm{mex}_1(u) \le 0$.

This case is never positive.

We add $m$ (filling the first gap) and remove $a$.

• If $a < m$: removing $a$ creates a gap at $a$, before the filled gap. Change $= a - m < 0$.
• If $a > m$: the set $\{0, \ldots, m\}$ is now complete. The new MEX is somewhere beyond $m$. This is the only positive case.

For the sub-case $a > m$, define $\mathrm{mex}_2(u) =$ the smallest value $\ge m + 1$ not in $S_u$. Equivalently, $\mathrm{mex}_2(u) = \mathrm{MEX}(S_u \cup \{m\})$. This can be computed during the same DFS as $\mathrm{mex}_1$ by querying max_right(mex1 + 1).

After adding $m$ and removing $a$ (where $a > m$):

• If $a < \mathrm{mex}_2(u)$: the gap at $a$ comes before $\mathrm{mex}_2(u)$. New MEX $= a$.
• If $a \ge \mathrm{mex}_2(u)$: the original second gap remains. New MEX $= \mathrm{mex}_2(u)$.

Combined: Change $= \min(a,\; \mathrm{mex}_2(u)) - m > 0$.

If $a < m$: every descendant contributes $\le 0$. Not worth it.

If $a > m$: positive contributions come from Case B descendants, negative from Case A descendants with $\mathrm{mex}_1(u) > a$.

So we only consider vertices $v$ with $p_v > \mathrm{mex}_1(v)$.

The gain formula:

We compute these two parts separately.

We need, for each vertex $v$ with $p_v = a$:

where $C = |\{u \in \mathrm{sub}(v) : \mathrm{mex}_1(u) \ge a\}|$ and $\Sigma = \sum_{\substack{u \in \mathrm{sub}(v) \\ \mathrm{mex}_1(u) \ge a}} \mathrm{mex}_1(u)$.

(Including $\mathrm{mex}_1(u) = a$ in the filter is fine since those contribute $0$.)

Sweep from $n - 1$ down to $0$. Maintain two segment trees indexed by Euler-tour position: one storing $\mathrm{mex}_1$ values, one storing counts.

• At value $i$: insert all vertices $u$ with $\mathrm{mex}_1(u) = i$.
• Then for the vertex $v$ with $p_v = i$: query the subtree range $[L_v, R_v)$ on both trees.
• $\text{neg}(v) = i \cdot C - \Sigma$.

This is $O(n \log n)$.

We need, for each $v$ with $a = p_v > m = \mathrm{mex}_1(v)$:

Both $v$ and the relevant descendants $u$ share the same $\mathrm{mex}_1$ value $m$. So group all vertices by $\mathrm{mex}_1$ and process each group independently.

Fix a group with $\mathrm{mex}_1 = m$. Each member $u$ has some $\mathrm{mex}_2(u)$ and some $p_u$. A member can act as an “operator” (if $p_u > m$) and also as a “beneficiary” to ancestors in the same group.

For operator $v$ with value $a = p_v$, beneficiary $u$ in $\mathrm{sub}(v)$ contributes:

Use a segment tree on the Euler tour. Initially, set each member $u$’s position to $\mathrm{mex}_2(u) - m$.

Process operators in decreasing order of $p_v$. Before processing operator $v$, transition all members $u$ with $\mathrm{mex}_2(u) > p_v$: change their entry from $(\mathrm{mex}_2(u) - m)$ to $(-m)$ and set a separate count to $1$.

Then for operator $v$:

Since we process in decreasing $p_v$, the transition threshold only decreases, so each member is transitioned at most once. Each group costs $O(k \log n)$ where $k$ is its size, totaling $O(n \log n)$.

After processing a group, clean up the segment tree entries (reset to $0$).

Each of the three phases (MEX computation, negative part, positive part) is $O(n \log n)$. Total: $O(n \log n)$ per test case.

• Code