HaskellでBTree書いてQuickCheckでテストしてみた

Haskellが得意なことのひとつに木構造の操作がありますね。
ということでBTree書いてみました。
初め二分木書こうと思っていたんだけど、あまりにも簡単すぎたのでB木に変更。
そしたら思ったより難しくなってしまい。


selectとか木を走査するだけの関数ならパターンマッチ使って一瞬で書けるんだけど、insertはそれに加えてちょっと厄介な木の変更が行われるので、どんな関数書けば良いかわからず結構悩んだ。
具体的に言うと、B木のinsertは次数(以下ではnodeDegree)を超えるデータが一つのノードに挿入されると、そのノードを分解して「上(根方向)のノードに影響を及ぼす」。
ということでinsert'をなんとかひねり出した。
BTreeの型に大分悩んだのだけど、もっと良い型あれば助言いただきたいです。
というのは、Nodeの情報(NodeElement)を全部リストにして押し込めたのだけど、これをしてしまうとNodeElementの要素Dat, Subが交互に来ることとか、リストの最初と最後はSubになるという特性(#とする)を「自分で担保しなければならない」。折角Haskellの型の表現力高いんだから型である程度なんとかしたいよねー、というところ。



単純のためキーはIntにして、パフォーマンスより分かり易さ重視で書いてみた。つもり。
あ、同じキーが来たときの動作とか知らないから置き換えてしまっている。
deleteはまだ書いてない。

data BTree a = Empty | Node [NodeElement a]
    deriving (Show, Eq)
-- BTreeのNodeの要素
data NodeElement a = Dat (Int, a) | Sub (BTree a)
    deriving (Show, Eq)
-- BTreeのNodeのSeek結果
data SeekResult a b = Found a | SubTree b
    deriving (Show, Eq)
-- BTreeのNodeにinsertした後の結果
data TreeState a b = Normal (BTree a) | Intermediate ([NodeElement b], NodeElement b, [NodeElement b])
    deriving (Show, Eq)

-- ノードの許容数
nodeDegree = 2

----------------------------------------------------------------------------------------------
-- 挿入
insert :: Int -> a -> BTree a -> BTree a
insert i d nod = finalize . insert' i d . Normal $ nod

insert' :: Int -> a -> TreeState a a -> TreeState a a
insert' i d (Normal Empty) = Normal (Node [Sub Empty, Dat (i,d), Sub Empty])
insert' i d (Normal (Node elems)) = case seekKey i elems of
    Found (l, _, r)                  -> Normal $ Node (l ++ [Dat (i,d)] ++ r)
    SubTree (l, Sub Empty, r)        -> normalize $ l ++ [Sub Empty, Dat (i,d), Sub Empty] ++ r
    SubTree (l, Sub nod@(Node _), r) -> normalize $ l ++ merge (insert' i d . Normal $ nod) ++ r
    _                                -> error "Invalid result of `seekKey'"
insert' i d int@(Intermediate (lElems, mid, rElems)) = insert' i d (Normal . Node . merge $ int)

seekKey i elems = seekKey' i (reverse elems) []
seekKey' _ [Sub n] acc = SubTree ([], Sub n, acc)
seekKey' i (Sub n:Dat (l,d):es) acc | l < i     = SubTree (reverse (Dat (l,d):es), Sub n, acc)
                                    | l == i    = Found (reverse es, Dat (l,d), Sub n:acc)
                                    | otherwise = seekKey' i es (Dat (l,d):Sub n:acc)
seekKey' _ _ _ = error "Invalid args"

normalize elems | length elems > (nodeDegree * 2 + 1) = toIntermediate elems
                | otherwise                  = Normal (Node elems)

toIntermediate elems | length elems < 3 = error "Length of arg invalid"
                     | otherwise       = let len = length elems
                                             n = if even nodeDegree then (len - 1) `div` 2 else (len - 1) `div` 2 - 1
                                             (left, rest) = splitAt n elems
                                             (mid, right) = splitAt 1 rest
                                         in Intermediate (left, head mid, right)

merge (Normal nod) = [Sub nod]
merge (Intermediate (lElems, mid, rElems)) = [Sub (Node lElems), mid, Sub (Node rElems)]

finalize (Normal n) = n
finalize (Intermediate (lElems, mid, rElems)) = Node [Sub (Node lElems), mid, Sub (Node rElems)]
        
-- foldl用
insert'' :: BTree a -> (Int, a) -> BTree a
insert'' t (i,d) = insert i d t

----------------------------------------------------------------------------------------------
-- 検索
select :: Int -> BTree a -> Maybe a
select i Empty = Nothing
select i (Node elems) = case seekKey i elems of
    Found (_, Dat (_,d), _)          -> Just d
    SubTree (_, Sub Empty, _)        -> Nothing
    SubTree (_, Sub nod@(Node _), _) -> select i nod
    _                                -> error "Invalid result of `seekKey'"

----------------------------------------------------------------------------------------------

main :: IO ()
main = do
    let tree = foldl insert'' Empty [(18,'a'), (5,'b'), (25,'c'), (33,'d'), (28,'e'), (10,'f'), (22,'g'), (13,'h'), (16,'i'), (7,'j')]
    print $ select 7 tree -- => Just 'j'
    print $ select 200 tree -- => Nothing
    print $ tree -- 後述


上記foldlによって生成した木は以下。参考URLと同順でキーを挿入している。PrettyPrintは今のところ手作業..

Node [
    Sub (Node [
        Sub (Node [
            Sub Empty,
            Dat (5,'b'),
            Sub Empty,
            Dat (7,'j'),
            Sub Empty
        ]),
        Dat (10,'f'),
        Sub (Node [
            Sub Empty,
            Dat (13,'h'),
            Sub Empty,
            Dat (16,'i'),
            Sub Empty
        ])
    ]),
    Dat (18,'a'),
    Sub (Node [
        Sub (Node [
            Sub Empty,
            Dat (22,'g'),
            Sub Empty,
            Dat (25,'c'),
            Sub Empty
        ]),
        Dat (28,'e'),
        Sub (Node [
            Sub Empty,
            Dat (33,'d'),
            Sub Empty
        ])
    ])
]

QuickCheckによるBTreeのテスト

QuickCheckを使ってこのBTreeの性質が正しいか確認してみよう。
QuickCheckは型の性質を利用して値を自動生成し、テストを行うものだ。
僕はBTreeという型を定義しているが、実際にはこのBTreeの型に当てはまる値ランダムに自動生成しても、その中にはBTreeでないものが大量に含まれてしまう。何故かと言えば、前述したように[NodeElement]には型に表されていない特殊な性質(#)を、僕自身が担保しているからだ。
実際にBTreeとなるのは、関数insertによってEmptyから成長した木のみだ。


ではどうするか。BTreeの性質を考えてみる。
BTreeは検索を高速にするために構成されたデータ構造だ。なぜ検索が高速なのか?キーがソートされているから。BTreeはSortedSetだ。キーの配置が少し複雑になっているだけで。
ならばBTreeからキーのリストに変換するヘルパー関数(toSortedKeyList)を書いてみよう。

toSortedKeyList :: BTree a -> [Int]
toSortedKeyList Empty = []
toSortedKeyList (Node []) = []
toSortedKeyList (Node (Sub n:rest)) = toSortedKeyList n ++ toSortedKeyList (Node rest)
toSortedKeyList (Node (Dat (i,_):rest)) = i : toSortedKeyList (Node rest)

すぐに書けた。これだけ小さければこの関数自体にバグが入る余地も少ないだろう。
さてこの関数でBTreeから[Int]に変換が出来る。変換後のリストはソートされているはずだ。
なので、任意の[Int]に対して、BTreeは次の性質が成立する。はず。

prop_btree :: [Int] -> Bool
prop_btree ns = toSortedKeyList (foldl' insert'' Empty $ zip ns [1..]) == sort ns

(fold' insert'' Empty $ zip ns [1..])の部分はBTreeの生成部分だ。データはBTreeの性質に関わらないため、ここでは[1..]などとしている。これをtoSortedKeyListによって変換したリストは、もとのリストをソートしたものに等しい(sort ns)。
試してみよう。

$ ghci 
GHCi, version 7.0.2: http://www.haskell.org/ghc/  :? for help
Loading package ghc-prim ... linking ... done.
Loading package integer-gmp ... linking ... done.
Loading package base ... linking ... done.
Loading package ffi-1.0 ... linking ... done.
Prelude> :load BTree.hs 
[1 of 1] Compiling Main             ( BTree.hs, interpreted )
Ok, modules loaded: Main.
*Main> quickCheck (prop_btree :: [Int] -> Bool)
Loading package extensible-exceptions-0.1.1.2 ... linking ... done.
Loading package old-locale-1.0.0.2 ... linking ... done.
Loading package time-1.2.0.3 ... linking ... done.
Loading package random-1.0.0.3 ... linking ... done.
Loading package array-0.3.0.2 ... linking ... done.
Loading package containers-0.4.0.0 ... linking ... done.
Loading package pretty-1.0.1.2 ... linking ... done.
Loading package template-haskell ... linking ... done.
Loading package QuickCheck-2.4.1.1 ... linking ... done.
*** Failed! Falsifiable (after 10 tests and 4 shrinks): 
[-21,-21]

あ、失敗した。 prop_btreeに[-21,-21]を適用したときには失敗してしまうようだ。
そういえば僕が作ったBTreeは同じキーをinsertすると古いものと置き換えてしまうのだった。
ではnubを使って重複をなくそう。

prop_btree' :: [Int] -> Bool
prop_btree' ns = toSortedKeyList (foldl' insert'' Empty $ zip ns [1..]) == (sort . nub $ ns)

reloadしてもう一度。

*Main> :reload
Ok, modules loaded: Main.
*Main> quickCheck (prop_btree' :: [Int] -> Bool)
+++ OK, passed 100 tests.
*Main> quickCheck (prop_btree' :: [Int] -> Bool)
+++ OK, passed 100 tests.
*Main> quickCheck (prop_btree' :: [Int] -> Bool)
+++ OK, passed 100 tests.

通った。何度実行してもちゃんとpassする。(バグっていても、生成するテストケースによってはたまに通ってしまうようだ)
僕が書いたBTreeはちゃんとBTreeしてくれているらしい。
最適化の余地は結構ありそうだけど。リストの代わりにVector使うとか、ノードの変更を破壊的にしてみるとか?


quickCheckの代わりにverboseCheck使うとテストケースが見えるようだ。

*Main> verboseCheck (prop_btree' :: [Int] -> Bool)
Passed:
[]
Passed:
[]
Passed:
[]
Passed:
[-1]
Passed:
[-1]
Passed:
[8,-2,-2,3]
Passed:
[6,6,1,0,-2]
Passed:
[11]
Passed:
[-17,-4,-15,23,30,-12,-22,-13]
Passed:
[]
Passed:
[-61,63,-61,-64,-34,-47,4]
Passed:

... (中略) ...

[-770465736879868345,-627449363952103343,-1535647965411000244,-1617518870131125585,-748790527772618410,-882579783187142686,-557598521175080646,-1224670409996353249,-1051025035886364551,-435915657678907530,1574244981112913865,-2162517402137240470,608246772695075684,627255218398136101,-2032405622871832601,-1959949395788039963,1141660388144951757,-2231756685795212628,-1164560351227784352,-181241041539842080,-422805536797723516,1105868777602430855,-121034860432626942,1010213561190766524,2207323701813970773,1824558824753279464,1659928916504648837,-356409826894394941,-349061123608670862,752482050516129743,1732210273013981767,-2001436668295935953,397001051167470637,-586306845267045940,1120474658668622134,1697868973446201684,1310255745061996168,1551497250807091072,-539107966158989349,-2075997266422967180,337635779612776167,-763570382140750685,642781941496758194,905319539862446088,-414222507919609479]
Passed:
[900494496553520522,184865048628130966,-1642873476856253220,1421902239354081659,195588487225024683]
Passed:
[21147071859887645,-3242210789371958852,2468093635359341973,-2639972716416303355,2136152609480382306,-2405479805855764766,-2295511696565431053,-3973867953844783728,4125895748378591025,-3323960712440986908,-495062995479804210,-272879800836578695,2646719957544972788,517206298251015434,4341048928782758814,-2842685228415062675,-1825303373649335299,-859434232805242681,-1818904926304787247,-3456077728945581090,-45448592099274405,1230178338494958757,3166297607921802246,1658724553043350174]
+++ OK, passed 100 tests.

なるほどね。