Haskellで二分ヒープ(優先度付きキュー)実装してみた

Haskellには、C++で言う所のpriority_queueにあたるモジュールがないので、Codeforcesでこれが必要な問題は解けないという問題を抱えていた（幸い、これまで自分が出てきたコンテストでそういう問題には当たらなかった、はずです）ので、自前で実装してみました。

実装したのは

型

Heap a

a型の変数を持つヒープ

関数

関数名	型	機能
make_heap	(a -> a -> Bool) -> [a] -> Int -> Heap a	順序付けする関数（後述）と要素のリスト、ヒープの最大サイズを受け取ってヒープを構築
pop_heap	Heap a -> Maybe (a, Heap a)	最大（最小）値を取り出す。ヒープが空の時はNothingを返す
push_heap	Heap a -> a -> Heap a	ヒープに要素を挿入

です。他にも細々とした補助関数は定義しましたが。

コードは以下。
なお、ヒープが本来要求される計算量
make : O(n)
pop : O(log n)
push : O(log n)
を達成するにはGHCで-Oあるいは-O2オプションをつけてコンパイルすることが必要です。

import Data.Array.ST
import Control.Monad.ST
import Data.Array
import Data.List
import Monad
data Heap a = H {size::Int,ary::Array Int a,pol:: a -> a -> Bool}
instance (Show a) => Show (Heap a) where show H {size=n,ary=a} = "Heap [" ++ (intercalate "," (map (show.(a!)) [1..n]))++"]"
build_heap :: (a -> a -> Bool) -> [a] -> Int -> Array Int a
build_heap policy l n =
    runSTArray $ do
      let len = length l
      ary <- newListArray (1,n) l
      let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei
      let heapity_rec i = let left = 2*i;right = 2*i+1 in
                          do elem <- readArray ary i
                             left_elem <- if left<=len then readArray ary left else return elem
                             right_elem <- if right<=len then readArray ary right else return elem
                             (if left<=len&&policy left_elem elem then
                                  if right<=len&&policy right_elem left_elem then
                                      swap i right elem right_elem>>heapity_rec right
                                  else
                                      swap i left elem left_elem>>heapity_rec left
                              else
                                  when (right<=len&&policy right_elem elem) (swap i right elem right_elem>>heapity_rec right))
      let beg = div len 2 in mapM_ heapity_rec [beg,beg-1..1]
      return ary
make_heap :: (a -> a -> Bool) -> [a] -> Int -> Heap a
make_heap policy l n = H {size=length l,ary=(build_heap policy l n),pol=policy}
top_heap :: Heap a -> a
top_heap = (!1).ary
pop_heap :: Heap a -> Maybe (a, Heap a)
pop_heap heap@H{size=len,ary=a,pol=policy} = if len==0 then Nothing else Just (top_heap heap, heap{size=len-1,ary=heapity})
    where
      heapity  =
          runSTArray $ do
            ary <- unsafeThaw a
            readArray ary len >>= writeArray ary 1
            let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei
            let heapity_rec i = let left = 2*i;right = 2*i+1 in
                                do elem <- readArray ary i
                                   left_elem <- if left<=len then readArray ary left else return elem
                                   right_elem <- if right<=len then readArray ary right else return elem
                                   (if left<=len&&policy left_elem elem then
                                        if right<=len&&policy right_elem left_elem then
                                            swap i right elem right_elem>>heapity_rec right
                                        else
                                            swap i left elem left_elem>>heapity_rec left
                                    else
                                        when (right<=len&&policy right_elem elem) (swap i right elem right_elem>>heapity_rec right))
            heapity_rec 1
            return ary
push_heap :: Heap a -> a -> Heap a
push_heap heap@H{size=len,ary=a,pol=policy} x = H{size=len+1,ary=heapity,pol=policy}
    where
      heapity =
          runSTArray $ do
            ary <- unsafeThaw a
            writeArray ary (len+1) x
            let swap i j ei ej = writeArray ary i ej >> writeArray ary j ei
            let heapity_rec i = do
                         let parent = div i 2
                         elem <- readArray ary i
                         parent_elem <- readArray ary parent
                         when (i>1&&policy elem parent_elem) (swap i parent elem parent_elem >> when (parent>1) (heapity_rec parent))
            when (len>0) (heapity_rec (len+1))
            return ary

実装はよく見る配列による実装です。
二分木型を作ってその上で実装することも考えたのですが、
資料の多さからとりあえずこの方法を取りました。
Heap型には本体となる配列とヒープ末尾の位置を持たせています。

max・minヒープどちらにでも対応するために、
make_heapには要素を順序付けるための関数（述語）をもたせる仕様とし、

make_heap (<) ……でmin-ヒープ
make_heap (>) ……でmax-ヒープ

をそれぞれ構成できます。

STモナドでの配列操作が増えた結果かなり煩雑なコードになってしまったのが惜しいところですが、速度はまあまあです。
これを優先度付きキューとして使った蟻本の問題でもそのうち解いてみようと思います。

それではﾉｼ