形式的に検証したセグメント木の実装を使って競技プログラミングの問題を解く

　本記事では、定理証明支援系 Coq 上で形式的に検証したセグメント木の実装を使って、実際に競技プログラミングの問題を解いてみます。これによって、実用に堪えるコードであっても形式検証を行えることが分かることでしょう。

　前回の記事では、競技プログラミングで広く用いられている形式でのセグメント木の実装の正当性を定理証明支援系 Coq を使って証明したのですが、実際にその実装を使って問題を解いたりはしておらず、本当にコンテストで使えるような実装を検証しているのか伝わりにくい点もありました。競技プログラミングに親しんでいる人であれば、ああ良くある実装を Coq で検証しているんだなとコードを読んで理解できたかもしれませんが、そうでない人には形式的検証のために効率を犠牲にしているのではないかと疑いを持ったことでしょう。

　本記事では、前回の記事で形式的に検証したセグメント木の実装を使って、実際に競技プログラミングの問題を解いてみます。これによって、実際にコンテストで使っても問題ない、効率的な実装の検証を行ったことが保証できるかと思います。

OCaml による実装の抽出

　本記事では Coq 上で形式検証を行ったセグメント木の実装を使って競技プログラミングの問題を解きたい訳ですが、Coq で書いたコードに対応したコンテストサイトは殆ど存在しないため、まず Coq による実装をそれに対応する OCaml のコードに変換します。

　OCaml のコードに変換と言っても難しいことではなく、Coq 自体がそう言った機能を提供しているため、我々はただ使うだけで構いません。詳細が気になる人は公式ドキュメントを参照して下さい。

　実際に OCaml のコードに変換してみましょう。まともな実装に変換するためには若干の小細工（ペアノの公理に基づいて定義された Coq 標準ライブラリの自然数型を OCaml の整数型と同じものとみなしたりだとか）が必要ですが、以下のようなコマンドを Coq 上で実行すると

Extraction "segtreeQueries.ml" product' upper_bound' lower_bound'

次のような OCaml のコードが得られます。

(** val product_rec' :
    ('a1 -> 'a1 -> 'a1) -> (int -> 'a1) -> int -> int -> int -> int -> 'a1 ->
    'a1 -> 'a1 **)

let rec product_rec' mul x x0 x1 x2 x3 x4 x5 =
  if ( <= ) x3 x2
  then mul x4 x5
  else product_rec' mul x ((fun n -> n / 2) x0) (( + ) x1 x0)
         ((fun n -> (n + 1) / 2) x2) ((fun n -> n / 2) x3)
         (if (fun n -> 0 < n mod 2) x2 then mul x4 (x (( + ) x1 x2)) else x4)
         (if (fun n -> 0 < n mod 2) x3
          then mul (x (( + ) x1 (pred x3))) x5
          else x5)

(** val product' :
    'a1 -> ('a1 -> 'a1 -> 'a1) -> (int -> 'a1) -> int -> int -> int -> 'a1 **)

let product' idm mul segtree m l r =
  product_rec' mul segtree m 0 l r idm idm

(** val upper_bound_rec' :
    ('a1 -> 'a1 -> 'a1) -> 'a1 pred -> (int -> 'a1) -> int -> int -> int ->
    int -> int -> 'a1 -> (int -> 'a1 -> 'a2) -> 'a2 **)

let rec upper_bound_rec' mul p segtree i m n l r p0 cont =
  (fun fO fS n -> if n=0 then fO () else fS (n-1))
    (fun _ -> cont l p0)
    (fun i0 ->
    if ( <= ) r l
    then cont l p0
    else let p' =
           if (fun n -> 0 < n mod 2) l
           then mul p0 (segtree (( + ) n l))
           else p0
         in
         if if (fun n -> 0 < n mod 2) l then not (p p') else false
         then cont l p0
         else upper_bound_rec' mul p segtree i0 ((fun n -> n / 2) m)
                (( + ) n m) ((fun n -> (n + 1) / 2) l) ((fun n -> n / 2) r)
                p' (fun k p1 ->
                if ( <= ) r ((fun n -> n + n) k)
                then cont ((fun n -> n + n) k) p1
                else let p'0 = mul p1 (segtree (( + ) n ((fun n -> n + n) k)))
                     in
                     if p p'0
                     then cont (succ ((fun n -> n + n) k)) p'0
                     else cont ((fun n -> n + n) k) p1))
    i

(** val upper_bound' :
    'a1 -> ('a1 -> 'a1 -> 'a1) -> 'a1 pred -> (int -> 'a1) -> int -> int ->
    int -> int * 'a1 **)

let upper_bound' idm mul p segtree m l r =
  upper_bound_rec' mul p segtree m m 0 l r idm (fun x x0 -> (x, x0))

(** val lower_bound_rec' :
    ('a1 -> 'a1 -> 'a1) -> 'a1 pred -> (int -> 'a1) -> int -> int -> int ->
    int -> int -> 'a1 -> (int -> 'a1 -> 'a2) -> 'a2 **)

let rec lower_bound_rec' mul p segtree i m n l r p0 cont =
  (fun fO fS n -> if n=0 then fO () else fS (n-1))
    (fun _ -> cont r p0)
    (fun i0 ->
    if ( <= ) r l
    then cont r p0
    else let p' =
           if (fun n -> 0 < n mod 2) r
           then mul (segtree (( + ) n (pred r))) p0
           else p0
         in
         if if (fun n -> 0 < n mod 2) r then not (p p') else false
         then cont r p0
         else lower_bound_rec' mul p segtree i0 ((fun n -> n / 2) m)
                (( + ) n m) ((fun n -> (n + 1) / 2) l) ((fun n -> n / 2) r)
                p' (fun k p1 ->
                if ( <= ) ((fun n -> n + n) k) l
                then cont ((fun n -> n + n) k) p1
                else let p'0 =
                       mul (segtree (( + ) n (pred ((fun n -> n + n) k)))) p1
                     in
                     if p p'0
                     then cont (pred ((fun n -> n + n) k)) p'0
                     else cont ((fun n -> n + n) k) p1))
    i

(** val lower_bound' :
    'a1 -> ('a1 -> 'a1 -> 'a1) -> 'a1 pred -> (int -> 'a1) -> int -> int ->
    int -> int * 'a1 **)

let lower_bound' idm mul p segtree m l r =
  lower_bound_rec' mul p segtree m m 0 l r idm (fun x x0 -> (x, x0))

　これで OCaml で書かれた、セグメント木上の二分探索や区間積の実装が得られた訳ですが、実際に競技プログラミングの問題を解くにはセグメント木の構築や更新の実装も必要になります。Coq 上で副作用のある計算を扱うのは面倒なのと、元々実装があまり複雑でないのもあって検証をサボっていましたが、ここで一つ真面目に向き合ってみましょう。

　まず、モノイド型 M.t を格納するセグメント木を表す型を定義します。前回の記事では配列でセグメント木を表現した場合の実装を検証した訳なので当然 OCaml の配列型を使うんですが、配列の要素数（すなわちセグメント木全体の頂点数）からセグメント木の葉の数を求めるのは大変なので、葉の数 leaves も含めたレコード型として表現します。

type t = { leaves : int; data : M.t array }

　次に、全ての要素が単位元 M.e で初期化された、長さ m のセグメント木を構築する関数 make m を定義します。単位元同士でモノイドの演算 M.op を適用しても単位元のままなので、単に単位元だけ入った配列を確保するだけでセグメント木の不変条件（子に相当する頂点は、2つの親が保持する値にモノイドの演算を適用して得られた値を保持する）を満たしてくれます。

let make m =
  assert (0 <= m);
  (* Bits.popcnt は引数を2進数表現した際に1になる桁の数を返す関数のつもり
     OCaml の標準ライブラリには無いので、後で自分で定義する *)
  { leaves = m; data = Array.make (2 * m - Bits.popcnt m) M.e }

ここで、配列に必要な要素数（すなわちセグメント木全体の頂点数）は >> を右論理シフトとすると m + (m >> 1) + (m >> 2) + ... で表せますが、m を2進数表記した際に1になる桁の数を popcnt(m) と書くと、実はこの要素数は 2m - popcnt(m) で計算できます。愚直に要素数を計算しても O(log m) で済みますが、popcnt(m) は O(log log m) で求めるアルゴリズムが知られているので、それを使えば更に高速に求められることでしょう。まぁ、popcnt(m) は高々 log m 程度の値しか取らないので、ちょっぴりメモリを贅沢に使って長さ 2m の配列を確保するのも良いかもしれません。

　最後に、与えられたセグメント木 segtree の i 番目の要素 x を f x で更新する関数 update segtree i f を定義しましょう。これは、i 番目の葉の保持する値に f を適用した後、その葉の先祖にあたる頂点についても更新を行うことで実装できます。

(* 指定した頂点とその祖先の保持する値を更新する人
   i : 今更新しようとしている頂点の世代が配列のどの添字から格納されているか
   j : 今更新しようとしている頂点の親世代が配列のどの添字から格納されているか
   k : 今更新しようとしている頂点が、その世代で何番目に位置しているか
   n : 今更新しようとしている頂点の世代にいくつ頂点が存在するか
   d : セグメント木が格納されている配列 *)
let rec update_ancestors i j k n d =
  if k < n then begin
    d.(i + k) <- M.op d.(j + 2 * k) d.(j + 2 * k + 1);
    update_ancestors (i + n) i (k / 2) (n / 2) d
  end
let update { leaves; data } i f =
  assert (0 <= i && i < leaves);
  data.(i) <- f data.(i);
  update_ancestors leaves 0 (i / 2) (leaves / 2) data