セグメント木の一次元配列による実装の検証

　本記事では、セグメント木と呼ばれるデータ構造を一次元配列を使って実装した際の、区間積を求めるアルゴリズムの正しさを定理証明支援系 Coq を用いて証明します。

　前回の記事ではセグメント木に関するいくつかの操作を Coq で検証したのですが、ある競技プログラマの人から、葉の数が２の冪に限らない場合にも対応した１次元配列によるセグメント木の実装についても検証して欲しいといった指摘がありました。と言うのも前回の記事では葉の数が２の冪になっていない場合の実装を簡単にするために、ジャグ配列の要領で世代ごとに配列を分けて頂点のデータを格納していたのです。

　一次元配列によるセグメント木を実装して、それに関する操作を直接証明することも可能ではあるのですが、それが果たして良い証明なのかは疑問が残ります。本質的には木であるものを配列に埋め込んでいるせいで、セグメント木についての本質的なロジックとは別に、埋め込みの妥当性についても同じ証明の中で示さなければならないからです。やはり、証明であっても関心は分離すべきでしょう。

　本記事では、前回同様に世代ごとで配列を分けたセグメント木の区間積を求めるアルゴリズムを検証し、一次元配列を用いた実装との対応を示すことで、間接的に後者の実装の正しさを証明します。このアプローチによって、セグメント木についての本質的なロジックと、配列の二分木への埋め込みの妥当性ついてを、別の補題の証明として整理する事ができるでしょう。

葉の数が２の冪に限らない場合の区間積の実装

　まずは今回の記事で正しさを検証する、セグメント木を一次元配列を使って実装した際の区間積を求めるアルゴリズムがどういったものか見ていきましょう。

　前回の記事で区間積を求めるアルゴリズムを解説した際も、葉の数が２の冪である必要こそあったものの、セグメント木を一次元配列を使って実装する例を示しました。ここで再掲しておきましょう。

(* モノイドの型 m、単位元 idm、その演算 mul が定義されているものとする *)

(* 平衡二分木を幅優先探索し、各頂点の保持するデータを並べたものでセグメント木を表現する *)
type segtree = m array

(* [product_rec lp rp l r segtree]: 添字 [l, r) に対応する頂点が保持しているデータの積に、
   左から lp を、右から rp を掛けたものを返す *)
let rec product_rec : m -> m -> int -> int -> segtree = fun lp rp l r segtree ->
   if r <= l
   then mul lp rp
   else
     product_rec
       (if l mod 2 = 0 then lp else mul lp segtree.(l))
       (if r mod 2 = 0 then rp else mul segtree.(r - 1) rp)
       ((l + 1) / 2) (r / 2) segtree

let product : int -> int -> segtree = fun l r segtree ->
  let n = Array.length segtree / 2 in
  (* 葉にあたる頂点は n から 2 * n - 1 まで添字順に並んでいる
     単位元を掛けても元の値のままなので、lp と rp には単位元 idm を入れておく *)
  product_rec idm idm (l + n) (r + n) segtree

　なぜこの実装で葉の数が２の冪である必要があるかと言えば、そのような二分木、つまり完全二分木の頂点を幅優先順に 1-based で配列に割り当てると、添字 i の頂点の子は添字 2i と添字 2i + 1 に対応する事実を使っているからです。ではなぜ完全二分木ならそのような性質が成り立つかと言えば、頂点数が葉の数 n を使って n + n / 2 + (n / 2) / 2 + ... + 1 = 2n - 1 で表せるので、1-based で（先頭に余計な要素を追加して）配列に格納すると、ある世代で一番左に位置する頂点の添字が m の時、その子世代で一番左に位置する頂点の添字が 2m になるからですね。その世代で一番左に位置する頂点の添字 m と相対位置 j を使って、ある頂点の添字 i = m + j が表される時、その次の世代で一番左に位置する頂点の添字が 2m になるため、結局子の添字は 2m + 2j = 2(m + j) = 2i と 2m + (2j + 1) = 2(m + j) + 1 = 2i + 1 になります。

　要素数が２の冪とは限らない場合、ある世代の頂点がどの添字から格納されているかは簡単には求められないので、再帰で区間積を求める際はその情報も引き回してやる必要があります。OCaml による実装例を以下に示します。

(* 頂点数からすぐには葉の数が求まらないので、各頂点のデータだけでなく葉の数も覚えておく *)
type segtree = { leaves : int; data : m array }

(* [product_rec' lp rp l r m n segtree]:
     葉の数が m それ以外の頂点数が n のセグメント木 segtree を受け取り、
     添字 [l, r) に対応する葉が保持しているデータの積に、
     左から lp を、右から rp を掛けたものを返す *)
let rec product_rec' : m -> m -> int -> int -> int -> int -> m array -> m =
  fun lp rp l r m n segtree ->
    if r <= l
    then mul lp rp
    else
      product_rec'
        (if l mod 2 = 0 then lp else mul lp segtree.(n + l))
        (if r mod 2 = 0 then rp else mul segtree.(n + r - 1) rp)
        (* 葉の親世代の頂点数は、葉の数 m を用いて m / 2 で表される
           従って葉だけでなくその親世代も除いた頂点数は、
           葉を除いた頂点数 n を使って n - m / 2 と書ける *)
        ((l + 1) / 2) (r / 2) (m / 2) (n - m / 2) segtree

let product : int -> int -> segtree -> m = fun l r { leaves; data } ->
  (* 単位元を掛けても元の値のままなので、lp と rp には単位元 idm を入れておく *)
  product_rec' idm idm l r leaves (Array.length data - leaves) data

Coq による検証

　それでは実際に、セグメント木を一次元配列を使って実装した際の区間積を求めるアルゴリズムの検証を行いましょう。冒頭にも書いた通り関心の分離を狙って、頂点を世代ごとに別の配列に格納する実装を使った区間積を求めるアルゴリズムをまず検証し、一次元配列を用いた実装との対応を示すことで、間接的に後者の実装の正しさを証明します。

ジャグ配列を用いたセグメント木の実装に対する区間積の検証

　まず前回同様、ジャグ配列の要領で世代ごとに別の配列に分けて頂点の情報を格納する実装での区間積の検証を行いましょう。

　前回同様とは言っても、自然数二つを受け取る関数によるジャグ配列の表現は言わば要素数無限の二次元配列のようなものなので、一次元配列への埋め込みを考えれば有限長であることを意識した表現に改めるべきでしょう。今回はリストのリストを使ってセグメント木を表現し、区間積を実装します。

(* leaves: 葉のデータのリスト
   rest_nodes: 葉以外の頂点を世代ごとに分けてリストにした、リストのリスト
   セグメント木の世代数は有限であるとする定義を採用したおかげで、停止性を自動判定してくれる *)
Fixpoint product_rec l r lp rp leaves rest_nodes :=
  (* 親世代の頂点のリストを取り出して再帰を行う関数 *)
  let product_rec' := fun l r lp rp =>
    if rest_nodes is parents :: rest_nodes'
    then product_rec l r lp rp parents rest_nodes'
    (* 「valid」 なセグメント木ならこの節には行かないはずだが、とりあえず適当なデータを返す *)
    else lp * rp in
  if r <= l
  then lp * rp
  else
    product_rec' (uphalf l) r./2
      (if odd l then lp * nth idm leaves l else lp)
      (if odd r then nth idm leaves r.-1 * rp else rp).

(* 前述の leaves と rest_nodes の組によって表現されたセグメント木に対する区間積の実装 *)
Definition product l r '(leaves, rest_nodes) :=
  product_rec l r idm idm leaves rest_nodes.

　ここで実装したアルゴリズムによって区間積を正しく求めるためには、引数で受け取ったリストのリストがセグメント木としての制約を満たしている必要があります。前回の定式化で言う table_accumulate という公理がそれに相当しているのですが、今回は有限長のリストを使っているので、以下のようにリストの長さについての性質も制約に含めなくてはなりません。

(* 前述の leaves と rest_nodes の組がセグメント木を表現できているか *)
Fixpoint valid_segtree leaves rest_nodes :=
  if rest_nodes is parents :: rest_nodes'
  then
    (* 一番上以外の世代 *)
    (* 親の要素数は子の要素数の半分である
       ./2 は余りを切り捨てる除算なので葉の要素数が2の倍数である必要はない *)
    (size parents == (size leaves)./2) &&
    (* 子の頂点が持っているデータの積を、親は持たなくてはならない *)
    [forall i : 'I_(size leaves)./2,
      nth idm parents i == nth idm leaves i.*2 * nth idm leaves i.*2.+1] &&
    valid_segtree parents rest_nodes'
  else
    (* 根の世代を一番上としても良いが、証明を簡単にするために番兵を入れておく *)
    leaves == [::].

　妥当なセグメント木がどういったものか表現できたので、前回と同じような流れで正当性を証明できます。しいて違う所を挙げるとすれば、各世代の頂点数は有限であるとする定義を採用したので、添字が範囲内に収まっていることも言わなければならない事でしょうか。

Lemma product_rec_correct : forall rest_nodes l r lp rp leaves,
  valid_segtree leaves rest_nodes ->
  r <= size leaves ->
  product_rec l r lp rp leaves rest_nodes
    = lp * (\big[mul/idm]_(l <= i < r) nth idm leaves i) * rp.
Proof.
  elim => /= [ | parents ? IH ] l r ? ? leaves
        => [ /eqP -> /= | /andP [ /andP [ /eqP Hsize /forallP Hacc ] ? ] ] Hbound;
    case (leqP r l) => [ ? | Hlr ].
  - by rewrite big_geq ?Monoid.mulm1.
  - by move: (leq_trans Hlr Hbound).
  - by rewrite big_geq ?Monoid.mulm1.
  - rewrite IH ?Hsize ?half_leq // left_segment_acc right_segment_acc
      (accumulated_product (nth idm parents) (nth idm leaves))
        => [ | ? /andP [ ? /((@leq_trans _ _ _)^~(half_leq Hbound)) Hbound' ] ].
    { rewrite double_half double_uphalf !Monoid.mulmA.
      congr (_ * _). rewrite -!Monoid.mulmA. congr (_ * _).
      rewrite -!big_cat_nat ?leq_addl ?leq_subr //.
      - by rewrite (@leq_trans l.+1) // (leq_add2r _ _ 1) leq_b1.
      - case (boolP (odd l)) => Hoddl.
        + case (boolP (odd r)) => Hoddr.
          * move: subn1 (leqP r.-1 l) Hoddl (Hoddr) => -> [ ] //.
            by rewrite -{1}ltnS prednK ?odd_gt0 => // /(conj Hlr) /andP /anti_leq <- /= ->.
          * by rewrite subn0.
        + apply /(@leq_trans r.-1).
          * by rewrite -ltnS prednK ?(leq_ltn_trans (leq0n l)).
          * by rewrite -subn1 leq_sub2l ?leq_b1. }
    by have /eqP := Hacc (Ordinal Hbound').
Qed.

Theorem product_correct l r leaves rest_nodes :
  valid_segtree leaves rest_nodes ->
  r <= size leaves ->
  product l r (leaves, rest_nodes) = \big[mul/idm]_(l <= i < r) nth idm leaves i.
Proof.
  rewrite /product => /product_rec_correct Hvalid /Hvalid ->.
  by rewrite Monoid.mulm1 Monoid.mul1m.
Qed.

一次元配列を用いた実装との対応の検証

　続いて、頂点を世代ごとに別の配列に格納するセグメント木の実装を使った区間積のアルゴリズムと、一次元配列を用いた実装を使った区間積のアルゴリズムの対応を示しましょう。

　前に OCaml の実装を示しながら解説したように、葉の数とそれ以外の頂点数を再帰で引き回すことで、一次元配列を使ったセグメント木の実装に対しても区間積を求められます。

(* 前回はイキっていたので Fix を使って関数の停止性を明示したが、
   実は停止性の明らかではない関数を定義するだけなら Function でもできる *)
Function product_rec' segtree m n l r lp rp {measure id r} :=
  if r <= l
  then lp * rp
  else
    product_rec' segtree m./2 (n - m./2) (uphalf l) r./2
      (if odd l then lp * segtree (n + l) else lp)
      (if odd r then segtree (n + r.-1) * rp else rp).
Proof.
  move => ? ? ? l r ? ? /negbT Hgt.
  apply /ltP.
  by rewrite -divn2 ltn_Pdiv // (@leq_ltn_trans l) // ltnNge Hgt.
Defined.

Definition product' segtree m n l r :=
  product_rec' segtree m n l r idm idm.

　区間積を求めるアルゴリズムの対応を示すのに際して、ジャグ配列の要領で定義されたセグメント木を普通の配列の要領で定義したセグメント木に変換する関数 encode を定義しましょう。と言っても根から順番に並べていくだけなんですが。

Fixpoint encode {A} (nodes : seq (seq A)) :=
  if nodes is leaves :: rest_nodes
  then encode rest_nodes ++ leaves
  else [::].

Eval simpl in (encode [:: [:: 4; 5; 6; 7]; [:: 2; 3]; [:: 1]; [::]]).
(*   = [:: 1; 2; 3; 4; 5; 6; 7]
     : seq nat *)

　準備が整ったので、区間積を求めるアルゴリズムの対応を証明します。

Lemma product_rec_correspondence : forall rest_nodes leaves segtree l r lp rp,
  valid_segtree leaves rest_nodes ->
  r <= size leaves ->
  (forall i,
    i < size (encode rest_nodes) + size leaves ->
    segtree i = nth idm (encode rest_nodes ++ leaves) i) ->
  product_rec l r lp rp leaves rest_nodes
    = product_rec' segtree (size leaves) (size (encode rest_nodes)) l r lp rp.
Proof.
  elim => /= [ | ? ? IH ] ? segtree l r ? ?
        => [ /eqP -> /= | /andP [ /andP [ /eqP Hsize ? ] ? ] ] Hbound Hsegtree;
    rewrite product_rec'_equation;
    case (leqP r l) => //= Hlr.
  - by move: (leq_trans Hlr Hbound).
  - rewrite !Hsegtree ?ltn_add2l.
    + rewrite nth_cat ltnNge leq_addr /=
              nth_cat ltnNge leq_addr /= !addKn size_cat Hsize addnK
              (IH _ segtree) ?Hsize ?half_leq => //= i Hi.
      by rewrite Hsegtree ?nth_cat ?size_cat Hsize ?(leq_trans Hi (leq_addr _ _)) ?Hi.
    + by rewrite prednK // (leq_ltn_trans (leq0n l)).
    + by rewrite (@leq_trans r).
Qed.

Lemma product_correspondence rest_nodes leaves segtree l r :
  valid_segtree leaves rest_nodes ->
  r <= size leaves ->
  (forall i,
    i < size (encode rest_nodes) + size leaves ->
    segtree i = nth idm (encode rest_nodes ++ leaves) i) ->
  product' segtree (size leaves) (size (encode rest_nodes)) l r 
    = product l r (leaves, rest_nodes).
Proof.
  by rewrite /product' /product => /product_rec_correspondence Hvalid /Hvalid Hbound /Hbound.
Qed.

　Coq が読める人であれば、なぜセグメント木によって区間積を求められるかのロジックについて、証明中で一切言及されていない事に気付くでしょう。ここでの証明は二分木を一次元配列に埋め込むことの妥当性に終始しており、関心の分離が実現できている事が分かります。

一次元配列を用いたセグ木の実装に対する区間積の正当性

　世代ごとで配列を分けたセグメント木の区間積を求めるアルゴリズムを検証し、一次元配列を用いた実装との対応も示せたので、間接的に後者の正しさを証明します。とは言ってもこの定理の証明は言わばグルーコードのようなもので、大した量ではないです。

Theorem product'_correct rest_nodes leaves segtree l r :
  valid_segtree leaves rest_nodes ->
  r <= size leaves ->
  (forall i,
    i < size (encode rest_nodes) + size leaves ->
    segtree i = nth idm (encode rest_nodes ++ leaves) i) ->
  product' segtree (size leaves) (size (encode rest_nodes)) l r 
    = \big[mul/idm]_(l <= i < r) nth idm leaves i.
Proof.
  move => ? ? ?.
  by rewrite product_correspondence // product_correct.
Qed.

まとめ

　本記事では、セグメント木と呼ばれるデータ構造を一次元配列を使って実装した際の、区間積を求めるアルゴリズムの正しさを定理証明支援系 Coq を用いて証明しました。今回 Coq で記述したプログラムや証明の全文は Gist にアップロードしてあるので、手元で動かしてみたい方があればそちらをご利用下さい。

　今回証明のアプローチとして、前回の記事と同様に世代ごとで配列を分けたセグメント木の区間積を求めるアルゴリズムを検証し、一次元配列を用いた実装との対応を示すことで、間接的に後者の実装の正しさを証明する手法を取りました。このアプローチによって、セグメント木についての本質的なロジックと、配列の二分木への埋め込みの妥当性ついてが別の補題の証明として整理され、全体として見通しの良い証明になったのではないでしょうか。

　この記事ではセグメント木上の二分探索については扱いませんでしたが、区間積と同様の手法で効率良く検証できるものと思われます。とはいえ、元が大変な証明なので結構な時間を拘束されそうですから、一旦棚上げして別のネタでブログを書くことになりそうです。