[ABC] Median of Medians(ABC 107[D])

スポンサーリンク

概要

問題文 →
公式解説 →
自分の提出 →

公式解説の計算量 $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」といえば、線形時間で中央値(あるいはその他の順序統計量)を計算する面白いアルゴリズムです()。知らない人はついでにぜひ学んでみてください。

そのアルゴリズムの名前から作問された問題なのかなー?

タイトルとURLをコピーしました