終端 \(r\) による場合分け

\(r \in [0, N[\) に対して \(f(r)\) を次のように定義します。

\[ f(r) = \sum _ { 0 \le l \lt r } \mathrm{mex} (A[l, r[) \]

このとき答えは \(\sum_ {r = 1} ^ N f(r)\) になります。

\(f(r)\) を式変形する

すると、公式解説同様の考察により、\(f(r)\) が次のように表せます。

\[ \begin{aligned} f(r) &= \sum _ {i = 0} ^ {r - 1} \mathrm{mex}(A[i, r[) \\ &= \sum _ {i = 0} ^ {r - 1} \sum_{x=0}^N \big [ [0, x[ \ \subseteq A[i, r[ \big]\\ &= \sum _ {x = 0} ^ N \min _ { 0 \le x' \le x } b_ {x'} \end{aligned} \]

(ただし \(b _ x = \min \left \lbrace i \in [0, n[ \ \middle | x \not \in A[i, n[ \right \rbrace\) であり、命題 \(P\) に対して \(P\) が真ならば \([P] = 1\)、偽ならば \([P] = 0\) であるとします。)

Historical minimum で表現する

\((b_x)_{x=0}^N\) の hisotrical (strict) minimum positions を \(0 = i_0 < i_1 < \dots < i_{j-1} < i_k\) すると、\(f(r)\) は次のように書けます。

\[ f(r) = \sum _ {j = 1} ^ k (i_j - i_{j-1}) \cdot b_{i_{j-1}} \]

\(b_x\) を管理して \(f(r)\) を計算していきましょう

\(f(r) \ (r \in [1, N])\) を \(r\) の大きい順に計算することを目指します。この順で見るとと \(b_x\) が時系列に関して単調減少になるため、償却更新箇所数 \(O(1)\) (実際には \(2\) 以下) のコストで historical minimum と \(f(r)\) を管理・更新することができます。

更新箇所の特定と更新は、いわゆる map (C++ の std::map, Rust の BTreeMap など) でできて、これにより全体 \(O(N \log N)\) 時間で計算できます。またより速い predecessor データ構造 (\(64\)-分木、van Emde Boas 木、Y-fast Trie など) で管理することで、\(O(N \log N / \log w)\) 時間、または \(O(N \log \log N)\) 時間で計算できます。

実装例

• BTreeMap (Rust 70 ms)
• 64分木 (Rust 16 ms)

use proconio::input;
use std::collections::BTreeMap;

fn main() {
    input! {
        n: usize,
        a: [usize; n],
    }
    let mut b = vec![0; n + 1];
    let mut prev = vec![usize::MAX; n];
    for (i, &a) in a.iter().enumerate() {
        prev[i] = b[a];
        b[a] = i + 1;
    }
    let mut last = 0;
    let mut current = 0;
    let mut historical_min = BTreeMap::new();
    historical_min.insert(last, b[last]);
    for (x, &y) in b.iter().enumerate() {
        if b[last] > y {
            historical_min.insert(x, y);
            current += (x - last) * b[last];
            last = x;
        }
    }
    let mut ans = current;
    for (i, &x) in a.iter().enumerate().rev() {
        let (&first_x, &first_y) = historical_min.range(..=x).next_back().unwrap();
        let mut last_x = first_x;
        let mut last_y = first_y;
        let y = prev[i];
        b[x] = y;
        if y < last_y {
            loop {
                let (&k, &v) = historical_min.range(x + 1..).next().unwrap();
                current -= (k - last_x) * last_y;
                last_x = k;
                last_y = v;
                if v <= y {
                    break;
                }
                historical_min.remove(&k).unwrap();
            }
            current += (x - first_x) * first_y;
            current += (last_x - x) * y;
            historical_min.insert(x, y);
        }
        ans += current;
    }
    println!("{ans}");
}