セグメント木の形式的検証
Photo by Todd Quackenbush on Unsplash
本記事では、セグメント木と呼ばれるデータ構造の諸操作の中でも実装が複雑な、区間積を求めるクエリと、適当な始点を与えて条件を満たす最大の区間を二分探索するクエリの正しさを定理証明支援系 Coq で証明します。
AtCoder 社が競技プログラミングのためのライブラリ AtCoder Library を配布したことで、コンテストでは頻出なのにそれまで実装の難しさから忌避されがちだったデータ構造やアルゴリズムも、気軽に使われるようになりました。今の水色コーダーはフローネットワークの最小費用流を使う問題さえ解いてしまうと言うのだから驚きです。
そんな便利ライブラリ AtCoder Library のクラスだとか関数だとかをブラックボックスとして使えば確かに問題を解くには困らないのですが、プログラマとしてのスキルアップのために競技プログラミングをやっているのだとしたら、やはりライブラリの実装しているデータ構造やアルゴリズムの本質を理解したい気持ちになることでしょう。
個人的な意見ですが、データ構造やアルゴリズムを理解するための最も確実な方法は、その性質を証明することだと思っています。真面目にプログラムの理解を試みる際は一行一行ソースコードを追って、何故その実装で良いのかのコメントを必要に応じて残していくことと思いますが、そのコメントを行間無く数学的に厳密に書き下したものこそが証明ではないでしょうか?
本記事では、競技プログラミングで頻繁に用いられるにも関わらず実装の複雑さからブラックボックスとして用いられがちなデータ構造、セグメント木に着目し、その諸操作の中でも実装が複雑な、区間積を求めるクエリと、適当な始点を与えて条件を満たす最大の区間を二分探索するクエリの正しさを、定理証明支援系 Coq を用いて証明します。
セグメント木
セグメント木とは、(主に 0 から N-1 の)整数で添字付けられたデータを保持し、与えられた範囲の添字のデータをモノイドの演算で畳み込んだ値を求めるクエリや、与えられた添字に対応するデータを置き換える操作を、共に O (log N) で行えるデータ構造です。似たような操作を実現する手法には Fenwick tree や累積和などがありますが、逆元の存在を仮定しなくて良いので最小値や最大値を求めたい場合にも使える、データの更新が高速といった長所がセグメント木にはあります。その分実装の簡単さが犠牲になっている訳ですが
セグメント木の具体的な実装ですが、葉にあたる頂点に各添字に対応するデータを保持し、親にあたる頂点には子の保持するデータの積を保持するような二分木を用います。言い換えれば葉にあたる頂点は長さ一の区間に、親にあたる頂点は子の扱う範囲を繋げた区間に、含まれる添字を持ったデータをモノイドの演算で畳み込んだ値をそれぞれ保持している訳です。
実装のイメージとして、与えられたリストと同じデータを保持するセグメント木を構築する関数 of_list
の OCaml による実装を示します。
(* モノイドの型 m、単位元 idm、その演算 mul が定義されているものとする *)
(* 平衡二分木を幅優先探索し、各頂点の保持するデータを並べたものでセグメント木を表現する
平衡二分木を配列に埋め込む際のありがちなテクニック *)
type segtree = m array
(* [of_list l]: リスト l と同じデータを保持するセグメント木を返す関数
リストの長さは2の冪でなくてはならない *)
let rec of_list : m list -> segtree = fun l ->
let n = List.length l in
(* 葉の数が n の平衡二分木の頂点数は n + n / 2 + (n / 2) / 2 ... = 2 * n - 1
1-based にするためにダミーの要素を1つ追加して配列を確保する *)
let segtree = Array.make (2 * n) idm in
(* 1-based で幅優先順に頂点のデータを配列に格納すると、
葉にあたる頂点は添字 n から始まる *)
List.iteri (fun i m -> segtree.(i + n) <- m) l;
for i = n - 1 downto 1 do
(* 1-based で幅優先順に頂点のデータを配列に格納すると、
添字 i の頂点の子は添字 2 * i と添字 2 * i + 1 になる *)
segtree.(i) <- mul segtree.(2 * i) segtree.(2 * i + 1)
done;
segtree
(* m = int, mul = min の場合の例 *)
# of_list [ 1; 2; 3; 4 ];;
- : int array list = [| 0; 1; 1; 3; 1; 2; 3; 4 |]
ここでは一般的なセグメント木の実装に倣い、二分木の表現として配列への埋め込みを採用していますが、木構造らしく代数的データ型や構造体とポインタ等を使っても実装できます。
セグメント木に対する一つ目の操作、与えられた範囲の添字のデータの積を求めるクエリの実装ですが、各頂点に格納されている区間積のうち、なるべく広い区間のものを選んで掛け合わせることで実現できます。OCaml による具体的な実装例を参考に、どのように実現すれば良いのか見ていきましょう。
(* [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
基本的なアイデアはコメントにも書いた通り、親にあたる頂点は2つの子のデータの積を持っているので、親世代で区間の積を求めれば子世代の半分の長さの区間だけ考えれば良く、これを繰り返せば O (log N) で処理できると言うことです。もっとも、区間 [l, r)
の始点 l
や終点 r
が奇数の場合、端にあたる添字 l
及び r - 1
の頂点が保持するデータは親世代では上手く処理できませんから、ここだけ個別に扱う必要があります。
次にセグメント木に対しての二つ目の操作、与えられた添字に対応するデータを更新する操作ですが、これはそのデータを保持する葉を更新した後に、その葉の祖先にあたる頂点についても更新すれば良いだけです。これも OCaml による実装例を見てみましょう。
(* 添字 i に対応する頂点と、その祖先にあたる頂点が保持するデータを更新する *)
let rec propagate : int -> segtree = fun i segtree ->
if 0 < i then begin
segtree.(i) <- mul segtree.(2 * i) segtree.(2 * i + 1);
propagate (i / 2)
end
(* [update i f segtree]: セグ木 segtree が保持する添字 i のデータ a.(i) を
f a.(i) で更新する *)
let update : int -> (m -> m) -> segtree = fun i f segtree ->
let n = Array.length segtree / 2 in
segtree.(i + n) <- f segtree.(i + n);
propagate ((i + n) / 2)
葉ノードは生のデータを保持しているのでこれを処理するのは当然ですが、その親の保持するデータも、そのまた親の保持するデータも、書き換えられた葉のデータに依存しているので更新する必要があります。とはいえ、木が適切にバランスされていればその高さは頂点数の log を取ったもので収まりますから、この操作も O(log N) で処理することができます。
また、あまり本質的な操作ではありませんが、区間の始点ないし終点を固定して、区間積が条件を満たすような最大の区間を求める二分探索も考えることができます。これは前述の区間積を求めるクエリを使っても O((log N)^2) で実装することはできるのですが、工夫すれば O(log N) で実装できることが知られています。本記事では始点を固定して最大の区間を求める二分探索だけ検証しますが、終点を固定するバージョンについても同様に検証できることでしょう。
定理証明支援系 Coq とは?
フランス国立情報学自動制御研究所で開発が進められている、数学の定理証明を支援するソフトウェアです。
証明なんて紙とペンさえあれば書けるのに、何故ソフトウェアの助けを借りる必要があるのかと思うかもしれませんが、自然言語で書かれた証明は往々にして間違うもので、正しさを担保するために査読の形で人的資源が費やされてきました。これは言わば動的型付き言語でプログラムを書いて、バグが無いか目で確認するようなものですが、最初から静的型付き言語で書いておけば手間が省けるのにと思うことでしょう。定理証明支援系は、数学の証明の世界にも型チェッカを導入したようなものと言えます。
定理証明支援系は Coq 以外にも Isabelle/HOL、Agda、HOL、Lean など色々ありますが、Coq の特色としてカリー・ハワード同型対応に基づいて設計されていることが挙げられます。
カリー・ハワード同型対応とは、数学の世界における命題及び証明と、プログラミング言語の世界における型及びプログラムの間に存在する、一対一の対応関係です。要するに数学の命題はプログラミング言語の型に、数学の証明はプログラムに対応していると言っていて、例えば A ならば B のような命題は A を受け取って B を返す関数の型に対応します。
Coq はこのカリー・ハワード同型対応に基づいて設計されていて、単なる定理証明支援系ではなく、述語論理も表現できるくらい強力な型システムを持った関数型言語と見ることもできます。プログラムは証明に対応するので Coq で記述できるし、そのプログラムの性質を Coq 上で証明することもできるのです。
セグメント木の諸操作の検証
セグメント木の定式化
セグメント木についての操作を Coq で検証するにあたって、まずセグメント木を定式化することから始めましょう。
初めに、セグメント木に載せるモノイドの存在を仮定します。単位元が idm
で演算が mul
ですね。一応 mul x y
を x * y
と書けるようにしておきます。
From mathcomp Require Import all_ssreflect.
Section Segtree.
Variable A : Type.
Variable idm : A.
Hypothesis mul : Monoid.law idm.
Local Notation "x * y" := (mul x y).
次にセグメント木の定式化ですが、この記事で示した OCaml による実装例同様、配列に二分木を埋め込む手法を採用します。とはいえ、一次元配列に埋め込んでしまうと完全二分木以外の場合に親子関係がややこしくなるので、世代ごとに別の配列にしたジャグ配列の気持ちで定式化します。
(* [table n i]: 根から n 遡った世代の、左から i 番目の頂点が保持するデータを返す関数 *)
Variable table : nat -> nat -> A.
(* 親にあたる頂点は、子にあたる頂点のデータの積を保持する *)
Hypothesis table_accumulate :
forall n i, table n.+1 i = table n i.*2 * table n i.*2.+1.
一般的な木の定式化と言えば連結な閉路を持たないグラフだと思いますが、本記事ではセグメント木の検証のために、実装に近い定義を採用しています。
ここで、区間積を実装した際の重要な直感を証明しておきましょう。親世代で区間積を求めれば、子世代の倍速で求められると言うものです。
Lemma upper_table_product n l : forall r,
\big[mul/idm]_(l <= i < r) table n.+1 i = \big[mul/idm]_(l.*2 <= i < r.*2) table n i.
Proof.
elim => [ | r IHr ].
- by rewrite double0 !big_geq // leq0n.
- case (leqP l r) => H.
+ by rewrite doubleS big_nat_recr // IHr table_accumulate Monoid.mulmA !big_nat_recr ?leq_Sdouble ?leq_double.
+ by rewrite !big_geq // leq_double.
Qed.
区間積を求める操作の検証
さて、セグメント木の定式化もできたことですし、区間積を求めるアルゴリズムを Coq で実装し、検証していきましょう。
区間積を Coq で実装する際に少し困るのは、その実装の再帰関数は Coq が停止性を自動判定するには複雑すぎる事です。停止しない関数というのはカリー・ハワード対応で言うと循環論法に相当するため、Coq では矛盾を防ぐために停止する関数しか許していないのです。
従って自分で停止性を証明しなければ関数の定義すらさせてもらえないのですが、折角プログラムとその停止性の証明を一緒に書くので、プログラムの正しさの証明も一緒に書いてしまいます。依存型まみれで少し読みづらいのですが、基本的にはコメントに書いた OCaml のコードと同じような事をしています。
(*
let rec product_rec r n l lp rp =
if r <= l
then lp * rp
else
product_rec (r / 2) (n + 1) ((l + 1) / 2)
(if l mod 2 = 0 then lp else lp * table n l)
(if r mod 2 = 0 then rp else table n (r - 1) * rp)
*)
Program Definition product_rec :=
Fix Wf_nat.lt_wf
(fun r => forall n l lp rp,
{ p | p = lp * (\big[mul/idm]_(l <= i < r) table n i) * rp })
(fun r product_rec n l lp rp =>
( if r <= l as b return (r <= l) = b -> _
then fun Hle => exist _ (lp * rp) _
else fun Hgt =>
let Hdecr := _ in
let (p, _) :=
product_rec r./2 Hdecr n.+1 (uphalf l)
(if odd l then lp * table n l else lp)
(if odd r then table n r.-1 * rp else rp) in
exist _ p _ ) erefl).
Coq ではこのように、Program
タクティックを用いて適当な証明項を _
で誤魔化しながらプログラムを書くと、証明項を対話的に補完するためのモードが立ち上がります。
まずはベースケースとなる r <= l
の場合の正しさを証明しなければなりませんね。そもそも区間の長さが 0 なので、区間積が正しく求められているのは自明です。
Next Obligation.
by rewrite big_geq // Monoid.mulm1.
Qed.
次にプログラムの停止性の証明です。引数 r
が単調に減少するので停止すると言っています。
Next Obligation.
apply /ltP.
by rewrite -divn2 ltn_Pdiv // (@leq_ltn_trans l) // ltnNge Hgt.
Qed.
最後に、l < r
の場合に区間積が求められていることの証明です。親世代で区間積を求めれば子世代の倍速で求められることは補題で証明しましたから、あとは if
式で書いた両端の部分と繋ぎ合わせると元の区間の積に一致する事を言えば良いです。
Next Obligation.
rewrite upper_table_product double_half double_uphalf left_segment_acc right_segment_acc !Monoid.mulmA.
congr (_ * _).
rewrite -!Monoid.mulmA.
congr (_ * _).
rewrite -!big_cat_nat ?leq_addl ?leq_subr //.
- apply /(@leq_trans l.+1).
+ by rewrite (leq_add2r l _ 1) leq_b1.
+ by rewrite ltnNge Hgt.
- case (@idP (odd r)) => Hoddr.
+ case (@idP (odd l)) => Hoddl.
{ rewrite subn1 ltn_predRL.
case (leqP l.+2 r) => //.
rewrite ltnS => Hleq.
have Hsucc : r = l.+1 by apply /anti_leq; rewrite Hleq ltnNge Hgt.
by move: Hsucc Hoddl Hoddr => -> /= ->. }
move: Hoddr Hgt => /odd_gt0 /prednK => {1}<- /negbT.
by rewrite -ltnNge ltnS subn1.
+ apply /(@leq_trans l.+1).
* by rewrite (leq_add2r l _ 1) leq_b1.
* by rewrite subn0 ltnNge Hgt.
Defined.
少し気を付けないといけないのは、親世代で処理できない両端の部分は子世代で処理するため、区間の長さは最大で2小さくなると言うことです。引数で与えられた区間 [l, r)
の長さが 1 しかない場合、余計な値まで掛けてしまわないか?と証明が進んでいくうちに Coq が尋ねてきますが、[l, r)
の長さが 1 の場合(つまり r = l + 1
の場合)、l
と r
のどちらかは偶数なので、そのような場合は両端を同時に掛けることがない訳です。
あとは引数 lp
及び rp
に単位元を渡して、使いやすい形の関数を得ても良いでしょう。これで区間積を求めるアルゴリズムの検証ができました。
(* let product = product_rec r 0 l idm idm *)
Program Definition product l r : { p | p = \big[mul/idm]_(l <= i < r) table 0 i } :=
let (p, _) := product_rec r 0 l idm idm in
exist _ p _.
Next Obligation.
by rewrite Monoid.mulm1 Monoid.mul1m.
Qed.
始点を固定して最大の区間を求める操作の検証
前述の通り、始点を固定した区間積の結果に応じて最大の区間を二分探索する操作も、Coq で実装し、検証します。ここでは実装の詳細には立ち入りませんが、親世代では倍速で探索できるので、左端と右端だけ線形探索して丸投げしている直感だけ得られれば理解できるかと思います。
(* モノイドの値を受け取る適当な述語 *)
Variable P : pred A.
(* P に [l, i) の区間積の値を入れると、i = m を境にして真偽が入れ替わる際、
その境界にあたる m と [l, m) の区間積を返す関数 *)
(*
let rec upper_bound_rec r n l lp =
if r <= l
then (l, lp)
else
let lp' = if l mod 2 = 0 then lp else lp * table n l in
if l mod 2 = 1 && not (P lp')
then (l, lp)
else
let (m, p) = upper_bound_rec (r / 2) (n + 1) ((l + 1) / 2) lp' in
if r <= m * 2
then (m * 2, p)
else
let p' = p * table n (m * 2) in
if P p'
then (m * 2 + 1, p')
else (m * 2, p)
*)
Program Definition upper_bound_rec :=
Fix Wf_nat.lt_wf
(fun r => forall n l lp,
(exists2 m, l <= m <= r &
forall k, P (lp * \big[mul/idm]_(l <= i < k) table n i) = (k <= m)) ->
{'(m, p) | [ /\ l <= m <= r,
p = lp * \big[mul/idm]_(l <= i < m) table n i &
forall k, P (lp * \big[mul/idm]_(l <= i < k) table n i) = (k <= m) ]})
(fun r upper_bound_rec n l lp Hex2 =>
( if r <= l as b return (r <= l) = b -> _
then fun Hle => exist _ (l, lp) _
else fun Hgt =>
let lp' := if odd l then lp * table n l else lp in
( if odd l && ~~ P lp' as b return odd l && ~~ P lp' = b -> _
then fun Htrue => exist _ (l, lp) _
else fun Hfalse =>
let (pair, _) := upper_bound_rec r./2 _ n.+1 (uphalf l) lp' _ in
( let (m, p) as pair0 return pair = pair0 -> _ := pair in fun _ =>
( if r <= m.*2 as b return (r <= m.*2) = b -> _
then fun Hge => exist _ (m.*2, p) _
else fun Hlt =>
let p' := p * table n m.*2 in
( if P p' as b return P p' = b -> _
then fun HP => exist _ (m.*2.+1, p') _
else fun HnP => exist _ (m.*2, p) _ ) erefl ) erefl ) erefl ) erefl ) erefl).
大変なんですが、これも頑張ったら正当性を証明できます。
(* r <= l の場合 *)
Next Obligation.
move: (@anti_leq l r) (H) H0 => <-.
+ move => /anti_leq ->.
by rewrite big_geq // Monoid.mulm1 !leqnn.
+ by move: H Hle => /andP [ Hle1 /(leq_trans Hle1) -> ] ->.
Qed.
(* 左端がはみ出ており、そこで真偽が切り替わる場合 *)
Next Obligation.
move: H Htrue (H0) => /andP [ Hle1 ? ] /andP [ -> ].
by rewrite leqnn (@leq_trans Hex2) // big_geq // Monoid.mulm1
-(@big_nat1 _ idm mul) H0 -ltnNge ltnS => /(conj Hle1) /andP /anti_leq <-.
Qed.
(* r が単調に減少することによる停止性の証明 *)
Next Obligation.
apply /ltP.
by rewrite -divn2 ltn_Pdiv // (@leq_ltn_trans l) // ltnNge Hgt.
Qed.
(* 子世代で真偽が切り替わる境界が存在するなら、親世代でもそれが存在することの証明 *)
Next Obligation.
exists Hex2./2 => [ | k ].
- move: H uphalf_half andbT (@idP (odd l)) (negbT Hfalse)
=> /andP [ ? /half_leq -> ] -> -> [ Hodd /= | ? ? ].
+ rewrite negbK -{1}(@big_nat1 _ idm mul) H0 => /half_leq.
by rewrite /= uphalf_half Hodd.
+ exact /half_leq.
- rewrite upper_table_product double_uphalf.
move: H (@idP (odd l)) Hfalse => /andP [ ? ? ] [ Hodd /= /negbT | ? ? ].
+ case (leqP k.*2 l) => [ Hle | ? ? ].
* by rewrite (big_geq (leqW Hle)) negbK Monoid.mulm1 -(doubleK k) half_leq // (@leq_trans l).
* by rewrite -Monoid.mulmA -(@big_nat1 _ idm mul) -big_cat_nat // H0 leq_double'.
+ by rewrite H0 leq_double'.
Qed.
(* 親世代で境界が右端にあった場合 *)
Next Obligation.
move: (H1) H => /andP [ ? Hle ] [ /andP [ ] ].
rewrite -(leq_double (uphalf l)) -(leq_double m) upper_table_product double_half double_uphalf left_segment_acc -Monoid.mulmA => ? ?.
rewrite -big_cat_nat ?leq_addl => // -> /(_ m).
rewrite upper_table_product double_uphalf -Monoid.mulmA -big_cat_nat ?leq_addl // leqnn => ?.
rewrite (@anti_leq m.*2 Hex2) // -H2.
by rewrite {2}(@anti_leq m.*2 r) ?Hle ?andbT // {1}(@leq_trans (r - odd r)) ?leq_subr.
Qed.
(* 区間の右端の余った部分に境界があった場合 *)
Next Obligation.
move: (H1) H HP => /andP [ ? ? ] [ /andP [ ] ].
rewrite -(leq_double (uphalf l)) -(leq_double m) upper_table_product double_half double_uphalf left_segment_acc => Hge' ?.
rewrite -Monoid.mulmA -big_cat_nat ?leq_addl => // -> /(_ Hex2./2).
rewrite upper_table_product double_uphalf double_half -!Monoid.mulmA => Hub.
rewrite -(@big_nat1 _ idm mul) -big_cat_nat ?(@leq_trans (odd l + l) l m.*2 (leq_addl _ _)) // H2 => Hgt'.
rewrite (@anti_leq m.*2.+1 Hex2) //.
move: Hub.
rewrite -big_cat_nat ?leq_addl //.
- rewrite H2 leq_subr -{1}leq_double double_half leq_subLR => /(@Logic.eq_sym _ _ _) H.
by move: (leq_add2r m.*2) (leq_b1 (odd Hex2)) andbT => <- /(leq_trans H) -> ->.
- move: (Hgt').
by rewrite -(@leq_subRL 1 m.*2 Hex2) ?(leq_ltn_trans (leq0n _) Hgt') => // /(leq_trans Hge') /((@leq_trans _ _ _)^~(leq_sub2l _ (leq_b1 _))).
Qed.
(* 区間の右端は余っていたけど、そこには境界は無かった場合 *)
Next Obligation.
move: H1 H (negbT HnP) => /andP [ Hge0 Hle0 ] [ /andP [ ] ].
rewrite -(leq_double (uphalf l)) -(leq_double m) upper_table_product double_half double_uphalf left_segment_acc => Hge' Hle'.
rewrite -Monoid.mulmA -big_cat_nat ?leq_addl => // -> /(_ m).
rewrite leqnn upper_table_product double_uphalf -!Monoid.mulmA => Hub.
rewrite -(@big_nat1 _ idm mul) -big_cat_nat ?(leq_trans (leq_addl _ _) Hge') //= H2 -leqNgt => Hle''.
rewrite (@anti_leq m.*2 Hex2) //.
move: Hub.
by rewrite -big_cat_nat ?leq_addl // H2 => ->.
Qed.
(* let upper_bound l r = upper_bound_rec r 0 l idm *)
Program Definition upper_bound l r
(Hex2: exists2 m, l <= m <= r &
forall k, P (\big[mul/idm]_(l <= i < k) table 0 i) = (k <= m)):
{'(m, p) | [ /\ l <= m <= r,
p = \big[mul/idm]_(l <= i < m) table 0 i &
forall k, P (\big[mul/idm]_(l <= i < k) table 0 i) = (k <= m) ]} :=
let (pair, _) := upper_bound_rec r 0 l idm _ in
exist _ pair _.
Next Obligation.
exists Hex2 => // ?.
by rewrite Monoid.mul1m.
Qed.
Next Obligation.
case H => ? -> Hub.
split => // [ | ? ].
- exact /Monoid.mul1m.
- by rewrite -Hub Monoid.mul1m.
Qed.
まとめ
本記事では、セグメント木と呼ばれるデータ構造に対する諸操作の正しさを、定理証明支援系 Coq を用いて数学的に証明し、バグが無いことを確認しました。今回 Coq で記述したプログラムや証明の全文は Gist にアップロードしてあるので、手元で動かしてみたい方があればそちらをご利用下さい。
リソースを滅茶苦茶使うので流石に業務では Coq で検証まではしませんが、このプログラムはどう動いているのか、厳密に考えてみる習慣はコードの品質を向上する助けになっているような気がします。僕の所属する DX チームには計算機科学の素養を持ったメンバーが所属しているので、興味があれば是非応募してみて下さい。