このエントリはML Advent Calendar 2015の14日目の記事です。 13日目の記事はよんたさんの関数型プログラミングにおける再帰関数の考え方でした。

ML Advent Calendar 2015を追っている方なら大体察しが付くと思いますが、この記事はobj_magicさんの A GADT implementation of AVL treeに刺激を受けて書きました。 GADTと型レベル自然数を使ってAVL木を書けるのなら、他のデータ構造も同じノリで実装してみると面白そうだと思ったので…まぁ二番煎じですね。

二項ヒープとは

マージが速いヒープの実装の一種らしいです。僕も二項ヒープを実装したのはこれが初めてで詳しい事は分からないので、気になった方は Amazon.co.jp： Purely Functional Data Structures: Chris Okasaki: 洋書みたいなデータ構造の本を読むと良いと思います。

型レベル自然数

二項ヒープを定義する前に型レベル自然数を定義します。zが0を表すつもりの型、n sがn + 1を表すつもりの型ですね。

type z = Z_
type 'n s = S_ : 'n -> 'n s

Z_とかS_は使わないから定義しなくても良いと思うかもしれませんが、これを定義しないと後々http://stackoverflow.com/questions/20692885/ と同じエラーが出ます。

自然数同士の比較をしたい場面が出てくるので型パラメーターと同じ値を持った自然数を定義します。

type _ nat =
  | Z : z nat
  | S : 'n nat -> 'n s nat

S (S (S Z))にz s s s nat型が付いて中々楽しいです。

あとはn <= mを表すつもりの(n, m) le型を定義して、簡単な補題を書いておきます。ペアノの公理っぽい形式の自然数同士の比較をするのは性能的にかなり微妙な気がしますが、まぁ要素数をNとするとせいぜいO(log N)同士の比較しか行わないので良いかと。

type (_, _) le =
  | LeEq : ('n, 'n) le
  | LeS : ('n, 'm) le -> ('n, 'm s) le

let rec le_n_s : type n m. (n, m) le -> (n s, m s) le = function
  | LeEq -> LeEq
  | LeS hle -> LeS (le_n_s hle)

let rec le_0_n : type n. n nat -> (z, n) le = function
  | Z -> LeEq
  | S n -> LeS (le_0_n n)

let rec le_ge_dec : type n m. n nat -> m nat -> [ `InL of (n, m) le | `InR of (m, n) le ] =
  fun n m ->
    match n, m with
    | Z, _ -> `InL (le_0_n m)
    | _, Z -> `InR (le_0_n n)
    | S n', S m' ->
        begin match le_ge_dec n' m' with
        | `InL hle -> `InL (le_n_s hle)
        | `InR hge -> `InR (le_n_s hge)
        end

データ構造の定義

型レベル自然数の定義ができたので、これを使って二項ヒープを表すデータ構造を作っていきましょう。

二項ヒープは二項木の集合として表されるので、二項ヒープを定義する前に二項木を定義しなければなりません。二項木とは、次数0の二項木は1つのノードからなり、次数kの二項木は次数0, 次数1, ... , 次数k-1の二項木を子に持つノードからなる、と言った具合で帰納的に定義されるデータ構造です。 Wikipediaの図を引用するとこんな感じですね。 f:id:fetburner:20151214064018p:plain

これに関しては実際のコードを見た方が早いでしょう。

type 'n tree = Node of elt * 'n tree_vector
and _ tree_vector =
  | TNil : z tree_vector
  | TCons : 'n tree * 'n tree_vector -> 'n s tree_vector

'n treeが次数nの二項木を表す型、'n tree_vectorが次数0からn-1までの二項木を集めたものを表す型です。ほげ_vectorみたいな型をいちいち手書きしたくはなかったですが、OCamlにはHigher-Kinded Polymorphismを扱うための機能が無いので…

GADTと型レベル自然数を持ち出したのは次数の情報を型にエンコードしたい、これに尽きます。次数の情報を型にエンコードしたので Node (1, TCons (Node (2, TNil), TCons (Node (3, TNil), TNil))) のように不正な二項木を表すプログラムがill-typedとなります。

二項木が定義できたので肝心の二項ヒープを定義していきます。

type _ tree_opt_vector =
  | TONil : z tree_opt_vector
  | TOCons : 'n tree option * 'n tree_opt_vector -> 'n s tree_opt_vector

type t = T : 'n nat * 'n tree_opt_vector -> t

'n tree_opt_vectorは次数n未満の二項木の集合を表しています。また、実際にライブラリを使う際にいちいち次数を型に書くのも面倒ですから、二項ヒープを表す型tではこれを隠しています。後々次数を比較する必要があるのでそれも覚えてたり。

データ構造の操作

データ構造が定義できたので、それを操作する関数を書いていきましょう。

手近なところでは、同じ次数を持つ二つの二項木をマージする関数はこんな感じです。

(* val merge_tree : 'a tree -> 'a tree -> 'a s tree = <fun> *)
let merge_tree (Node (e1, ts1) as t1) (Node (e2, ts2) as t2) =
  if compare e1 e2 < 0 then Node (e1, TCons (t2, ts1))
  else Node (e2, TCons (t1, ts2))

次数が1増える事が型にエンコードできていて大変素晴らしいですね。

これを使うと同じ次数を上限とする二項木の集合をマージする関数が書けます。

(*val merge_tree_opt_vector :
  'n tree_opt_vector ->
  'n tree_opt_vector -> 'n tree_opt_vector * 'n tree option = <fun> *)
let rec merge_tree_opt_vector : type n. n tree_opt_vector -> n tree_opt_vector -> n tree_opt_vector * n tree option = fun tov1 tov2 ->
  match tov1, tov2 with
  | TONil, TONil -> (TONil, None)
  | TOCons (to1, tov1'), TOCons (to2, tov2') ->
      let (tov12, to3) = merge_tree_opt_vector tov1' tov2' in
      begin match to1, to2, to3 with
      | _, None, None -> (TOCons (to1, tov12), None)
      | None, _, None -> (TOCons (to2, tov12), None)
      | None, None, _ -> (TOCons (to3, tov12), None)
      | _, Some t2, Some t3 -> (TOCons (to1, tov12), Some (merge_tree t2 t3))
      | Some t1, _, Some t3 -> (TOCons (to2, tov12), Some (merge_tree t1 t3))
      | Some t1, Some t2, _ -> (TOCons (to3, tov12), Some (merge_tree t1 t2))
      end

もちろん世の中にあるのは同じ次数を上限とした二項木の集合ばかりではありませんから、二項ヒープをマージする際には次数を比較して高い方に合わせています。

(* val padding : 'n tree_opt_vector -> ('n, 'm) le -> 'm tree_opt_vector = <fun> *)
let rec padding : type n m. n tree_opt_vector -> (n, m) le -> m tree_opt_vector = fun t hle ->
  match hle with
  | LeEq -> t
  | LeS hle' -> TOCons (None, padding t hle')

(* t -> t -> t *)
let merge (T (n1, tov1)) (T (n2, tov2)) =
  match le_ge_dec n1 n2 with
  | `InL hle ->
      begin match merge_tree_opt_vector (padding tov1 hle) tov2 with
      | tov', None -> T (n2, tov')
      | tov', (Some _ as to0) -> T (S n2, TOCons (to0, tov'))
      end
  | `InR hgt ->
      begin match merge_tree_opt_vector tov1 (padding tov2 hgt) with
      | tov', None -> T (n1, tov')
      | tov', (Some _ as to0) -> T (S n1, TOCons (to0, tov'))
      end

ヒープなので最小値の検索と削除をやりたいところですが、これらの処理は共通点が多いので一緒に書いてしまいます。しかし最小値の削除は二分木の再構築を行うため重そうに思えたので、遅延評価っぽい事をして計算量を削減したつもりです。

(* val delete_min_tree_opt_vector_aux1 :
  'a tree_opt_vector ->
  'a tree option -> (elt * (unit -> 'a s tree_opt_vector)) option = <fun> *)
let delete_min_tree_opt_vector_aux1 tov = function
  | None -> None
  | Some (Node (e, tv)) -> 
      Some (e, fun () ->
        let (tov', to0) = merge_tree_opt_vector tov (tree_opt_vector_of_tree_vector tv) in
        TOCons (to0, tov'))

(* val delete_min_tree_opt_vector_aux2 :
  'a tree option ->
  ('b * (unit -> 'a tree_opt_vector)) option ->
  ('b * (unit -> 'a s tree_opt_vector)) option = <fun> *)
let delete_min_tree_opt_vector_aux2 to0 = function
  | None -> None
  | Some (e, tov) -> Some (e, fun () -> TOCons (to0, tov ()))

(*val delete_min_tree_opt_vector :
  'n tree_opt_vector -> (elt * (unit -> 'n tree_opt_vector)) option = <fun> *)
let rec delete_min_tree_opt_vector : type n. n tree_opt_vector -> (elt * (unit -> n tree_opt_vector)) option = function
  | TONil -> None
  | TOCons (to0, tov) ->
      begin match delete_min_tree_opt_vector_aux1 tov to0, delete_min_tree_opt_vector_aux2 to0 (delete_min_tree_opt_vector tov) with
      | None, None -> None
      | (Some _ as result1), None -> result1
      | None, (Some _ as result2) -> result2
      | (Some (e1, _) as result1), (Some (e2, _) as result2) ->
          if compare e1 e2 < 0 then result1
          else result2
      end

(* val find_min : t -> elt option = <fun> *)
let find_min (T (_, tov)) =
  match delete_min_tree_opt_vector tov with
  | None -> None
  | Some (e, _) -> Some e

(* val delete_min : t -> t option = <fun> *)
let delete_min (T (n, tov)) =
  match delete_min_tree_opt_vector tov with
  | None -> None
  | Some (_, tov') ->
      begin match n, tov' () with
      | S n', TOCons (None, tov'') -> Some (T (n', tov''))
      | _, tov'' -> Some (T (n, tov''))
      end

なぜOCamlで標準にサポートされている遅延評価の機能を使わなかったかと言うと、僕が単に使った事が無くて使い方が分からなかったためです。

ソースコード

今回書いたプログラムの全文はMisc/heap.ml at master · fetburner/Misc · GitHubに置いてあります。

fetburner.core

コアダンプ

GADTを用いた二項ヒープの実装

二項ヒープとは

型レベル自然数

データ構造の定義

データ構造の操作

ソースコード