概要
公式解説の計算量 $O(N\log N\log A)$ よりも良い、$O(N\log N)$ 時間の解法について説明します。
解法（概要）
基本方針は公式解説と同じです。結局、二分探索内で次を解くことが課題となります（ここまでは、公式解説を見てください）。
$A_i\in \{\pm 1\}$ であるような数列 $A$ が与えられる。$\sum_{l\leq i\leq r}A_i \geq 0$ となるような組 $(l,r)$ を数え上げよ。
累積和をとることで、次に帰着されます。
$B_0 = 0$ かつ $|B_i – B_{i+1}| = 1$ を満たすような数列 $B_i$ が与えられる。$l < r$ かつ $B_l \leq B_r$ となる組 $(l,r)$ を数え上げよ。
これは数列 $B$ の転倒数と呼ばれ、
- マージソートしながら $O(N\log N)$ 時間で計算
- Fenwick Tree を用いて $O(N\log N)$ 時間で計算
などの典型があります。この記事では、この転倒数の計算を、$|B_i – B_{i+1}| = 1$ に注目して $O(N)$ 時間で計算することにします。
組 $(l,r)$ を、$r$ ごとに数え上げます。次のようにすればよいです。
$B_i$ の度数分布集計テーブル $\text{count}$ を持って、次のように計算します。
- $\text{count}[0] = 1$、他は $0$ と初期化しておく
- $i = 1, 2, \ldots,$ 順に次を行う。
- 答に $\sum_{x\leq B_i} \text{count}[x]$ を加える（$r = i$ となる $(l,r)$ の数え上げ）
- $\text{count}[B_i]$ に $1$ を加える
$\text{count}[x]$ を更新しながら累積和 $\sum_{x\leq B_i} \text{count}[x]$ を取得する必要があり、Fenwick Tree を使うとこれはクエリごとに $O(\log N)$ 時間でできます。
累積和をとるべき範囲がクエリごとに $\pm 1$ しか変化しないことを踏まえると、これは $O(N)$ 時間で行うことができます。各時点で、$\text{cumsum}:=\sum_{x\leq B_i}\text{count}[x]$ を保持するようにします。$i$ が $1$ 増えると、$\text{count}[x]$ の値をひとつ増減させるだけで次に必要な $\text{cumsum}$ の値が手に入り、上の手続きは $O(N)$ 時間で行うことができます。
二分探索の判定問題が $O(N)$ 時間で解けたため、全体で $O(N\log A)$ 時間で解くことができます。細かいですが、座標圧縮して $A$ の要素だけを対象とした二分探索を行うと、$O(N\log N)$ 時間になります。
コメント
この問題、ABC-D を埋めた当時の私には、ABC-D で最難問題のひとつに感じていましたが、今見返すと、一瞬で反応できる典型を数個組み合わせるだけで簡単ですね。やり込み量が反映されやすい問題でしょうか。
あまり関係のない話ですが、「Median of Medians」といえば、線形時間で中央値（あるいはその他の順序統計量）を計算する面白いアルゴリズムです（■）。知らない人はついでにぜひ学んでみてください。
そのアルゴリズムの名前から作問された問題なのかなー？
