Editorial : A2. Lost Civilization (Hard Version)

The algorithm starts from an input of length $m$ and repeatedly inserts $x_i + 1$ immediately after $x_i$. So every inserted element is exactly 1 greater than the element to its left. The function $f(b)$ asks: what is the minimum input length needed to produce $b$?

Stack-based computation of $f$: Process $b$ from right to left, maintaining a stack.

• For each element $b_i$, while the stack top equals $b_i + 1$, pop it (this element was “generated” from $b_i$).
• Push $b_i$.

The final stack size is $f(b)$.

Why this works: If $b_i$ is followed (eventually, in the stack sense) by $b_i + 1$, then $b_i + 1$ could have been inserted by the algorithm. The stack greedily absorbs such elements. What remains on the stack are the “root” elements that must have been in the original input.

Example: $b = [3, 4, 5]$. Processing right to left: push 5, then 4 pops 5 and is pushed, then 3 pops 4 and is pushed. Stack = $[3]$, so $f = 1$. Indeed, starting from $[3]$, insert 4 after 3, then insert 5 after 4, producing $[3, 4, 5]$.

Example: $b = [1, 3, 5]$. No consecutive-difference-1 pairs, so nothing gets popped. Stack = $[5, 3, 1]$, $f = 3$.

We want $\sum_{l} \sum_{r \ge l} f(a[l..r])$.

Fix the right endpoint $r$ and sweep $l$ from $r$ down to 0 (i.e. extend the subarray to the left one element at a time). Maintain the stack incrementally: when we prepend $a_l$, apply the pop/push rule. After each step, the stack size is $f(a[l..r])$. Accumulate the stack sizes.

for (int right = n - 1; right >= 0; right--) {
    stack<int> st;
    long long contribution = 0;
    for (int i = right; i >= 0; i--) {
        while (!st.empty() && st.top() == a[i] + 1)
            st.pop();
        st.push(a[i]);
        contribution += st.size();  // f(a[i..right])
    }
    ans += contribution;
}

This is $O(n^2)$ because for each of the $n$ right endpoints, the inner loop does $O(n)$ work.

Instead of fixing $r$ and sweeping $l$, fix $l$ and ask: what is $\text{res}(l) = \sum_{r=l}^{n-1} f(a[l..r])$?

The total answer is $\sum_l \text{res}(l)$.

Now consider what happens when we move from left endpoint $l = i+1$ to $l = i$. Every subarray $a[i+1..r]$ becomes $a[i..r]$ (we prepend $a_i$), plus there is one new subarray $a[i..i]$.

For each individual right endpoint $r$, prepending $a_i$ to the stack:

• Pops: While the stack top equals $a_i + 1$, pop it. Each pop decreases $f$ by 1.
• Push: Push $a_i$. This increases $f$ by 1.

So $f(a[i..r]) = f(a[i+1..r]) + 1 - (\text{number of pops for right endpoint } r)$.

If we could track the total number of pops across all right endpoints, we could update res in one shot:

Simplifying (since $1 + (n - 1 - i) = n - i$):

The challenge is computing “total pops” efficiently.

Maintain a single “merged” stack of pairs $(\text{value}, \text{weight})$. The merged stack simultaneously represents the stacks for all current right endpoints.

What the weight means: Each entry’s weight is the number of right endpoints for which that entry exists in their individual stack. Equivalently, the weight is that entry’s contribution to res.

Invariant: res = sum of all weights in the merged stack.

Structural picture. Suppose at left endpoint $i$, the merged stack (bottom to top) is $[(v_1, w_1), (v_2, w_2), \ldots, (v_k, w_k)]$. The individual stacks look like:

The weights are non-decreasing from bottom to top: $w_1 \le w_2 \le \cdots \le w_k$. The topmost entry is shared by all right endpoints (weight $= n - i$, the total count).

When we process position $i$:

1. Pop absorbed entries. While the merged stack’s top has value $a_i + 1$:

   • The weight $w$ tells us how many right endpoints have $a_i + 1$ on top of their individual stacks. For all of them, $a_i$ absorbs $a_i + 1$, reducing each $f$ by 1.
   • Subtract $w$ from res. Pop the entry.
   • If the new top also has value $a_i + 1$, it means some right endpoints had a second $a_i + 1$ on their stack. Repeat.

2. Push the new element. Push $(a_i, n - i)$ onto the merged stack. The weight is $n - i$ because:

   • $n - 1 - i$ existing right endpoints each gain $a_i$ on top of their stack ($f$ increases by 1 each).
   • 1 new right endpoint $r = i$ is introduced, with $f(a[i..i]) = 1$.
   • Total: $n - i$ added to res.

3. Accumulate. Add res to the answer. At this point, $\text{res} = \sum_{r=i}^{n-1} f(a[i..r])$, which is the total contribution of all subarrays with left endpoint $i$.

At each iteration $i$, ans += res adds the sum of $f$ over all subarrays that start at position $i$:

So the iteration sweeps left endpoints from right to left. After all iterations, ans holds $\sum_l \sum_{r \ge l} f(a[l..r])$, the desired answer.

$a = [1, 2, 5, 6, 5]$, $n = 5$.

Answer: 25. ✓

At step $i = 2$: popping (6, 2) means both right endpoints $r = 3$ and $r = 4$ had 6 on top of their stack. Since $a_2 = 5$ and $5 + 1 = 6$, element 6 is absorbed. res drops by 2 (the weight), then increases by 3 (the push weight $n - i = 3$). Net: $+1$.

At step $i = 0$: popping (2, 4) means all 4 existing right endpoints had 2 on top. Since $a_0 = 1$ and $1 + 1 = 2$, element 2 is absorbed. res drops by 4, then increases by 5. Net: $+1$.

void solve() {
    int n;
    std::cin >> n;

    std::vector<int> a(n);
    for (int i = 0; i < n; i++) {
        std::cin >> a[i];
    }

    std::vector<std::pair<int, int>> stk;  // (value, weight)
    i64 ans = 0;
    i64 res = 0;  // sum of f(a[i..r]) for all r >= i
    for (int i = n - 1; i >= 0; i--) {
        // Pop entries absorbed by a[i]
        while (!stk.empty() && stk.back().first == a[i] + 1) {
            res -= stk.back().second;
            stk.pop_back();
        }
        // Push a[i] for all (n - i) right endpoints
        stk.emplace_back(a[i], n - i);
        res += n - i;
        // res now equals sum of f(a[i..r]) for r = i..n-1
        ans += res;
    }

    std::cout << ans << "\n";
}

Time complexity: $O(n)$. Each element is pushed and popped at most once.

Space complexity: $O(n)$ for the stack.