Avant-Garde Code

アバンギャルド・コード is 前衛的算譜.

第6回 アルゴリズム実技検定 M 解説

第6回 アルゴリズム実技検定 (PAST) の問題 M - 等しい数 の解説です。
この記事では、範囲更新クエリ (RUQ) または平衡二分探索木を利用した解法を紹介します。
問題文はリンク先を参照してください。 atcoder.jp

解法の全体構成

まず、解法の全体構成は次のようになります。

  1. 同じ値が連続している部分を1つの区間と考える。クエリによる値の変更を何らかのデータ構造で管理することにより、各クエリに対して「どの値が何個増えたか (減ったか)」を求める。 なお、1回のクエリにより、複数の値 (区間) が影響を受ける可能性がある。
  2. それぞれの値 (区間) に対する差分を全体の組合せ数に反映させる。

後半の 2. のほうは、ある値の個数を c としたときの組合せ数が {}_c C_2 = \frac{1}{2} c(c-1) であることから、各値に対する増減量が与えられれば、差分を逐次計算することで出力する値を O(1) で調整できます。

さて、1. については、各クエリに対し、関連する区間を削除および追加することを考えます。両端においても、区間が少しでも重なっていれば操作の対象とします。
これにより 2. がサブルーチンとして実行される回数を考えると、初期状態から存在する区間のうち削除されるものはたかだか N 個で、1回のクエリにより追加される区間はたかだか 3 個であることから、全体では O(N+Q) であることに気付いておく必要があります。

RUQ の利用

「同じ値が連続する区間のデータをどのように管理するか」がこの問題の難しい部分だと思います。 ここでは、範囲更新クエリ (Range Update Query; RUQ) を利用することを考えます。

RUQ は、1回のクエリで、数列の区間 (連続する複数のインデックス) に対して同一の値を設定することを指します。 具体的にはセグメント木 (Segment Tree) で実装すればよいです。
参照: Range Update Query (AOJ)

この問題では、任意のインデックスに対して、それが属する区間の情報を取得できればよさそう、というのがポイントです。 区間の情報というのは「左端と右端のインデックス」でも「左端のインデックスと長さ」でもよいです。

解法として、次の2通りを考えました。

(1) RUQ で、各インデックスが属する区間の情報を保持する。すなわち、ある区間が半開区間 [L, R) で表せるとき、それに含まれる全てのインデックスに対して、値のペア (L, R) を設定する。 提出コード (C#)

(2) 大きな方針は (1) と同様で、「各インデックスが属する区間の左端のインデックスを保持する RUQ」と「左端のインデックスをキーとして、右端のインデックス (開区間) を保持する配列」を用意する。 提出コード (C#)

実装は (1) のほうが少しシンプルだと思いますが、隣の区間の情報を探索するために毎回 RUQ を呼び出すことになり、定数倍が重くなります。 (2) であれば、隣の区間はすぐに得られます。

いずれの解法においても、RUQ をセグメント木により実装することで、時間計算量 O((N+Q) \log N) で実行できます。

平衡二分探索木の利用

上記と同様の考え方により、平衡二分探索木を利用して解くこともできます。 1つの区間を1つの要素として格納すればよいため、発想としてはこちらのほうが自然かもしれません。

なお、今回の問題を解くためには、複数の区間についての探索ができる平衡二分探索木である必要があります。
参照: Set: Range Search (AOJ)
参照: Map: Range Search (AOJ)

解法として、次の2通りを考えました。

(1) 順序付き集合 (set) を表す平衡二分探索木で、各区間の情報を保持する。 提出コード (C#)

(2) 順序付き連想配列 (map) を表す平衡二分探索木で、各区間の情報を保持する。 すなわち、左端のインデックスをキーとして、右端のインデックス (開区間) を値として保持する。 提出コード (C#)

いずれの解法においても、時間計算量 O((N+Q) \log N) で実行できます。

まとめ

通常の RUQ であれば、数列の値 A_i をそのまま保持させることを考えるでしょう。 しかし今回のように、区間に関する別の情報を保持することで応用範囲が広がることがあります。

補記

平方分割で実装してみるとどうしても不正解になってしまったため、別解を考えた結果です。