でかいチーズをベーグルする

でかいチーズはベーグルすべきです。

HaskellでB-treeを実装

すごいHaskell本を読んだのでなにか練習したいなと思っていたところ、某会*1で「B-treeなんて誰でも簡単に実装できますよね」と煽られたので実装してみた。すごく直感的に書けたので Haskell すごいなーと思った。実装方針は "Introduction to Algorithms" に書いてあるとおり。

すごいHaskellたのしく学ぼう!

すごいHaskellたのしく学ぼう!

アルゴリズムイントロダクション 第3版 総合版 (世界標準MIT教科書)

アルゴリズムイントロダクション 第3版 総合版 (世界標準MIT教科書)

さっそく実装

B-treeを知っている人を対象にして書きます。B-treeを知らない人は Introduction to Algorithms を読むなりWebでいい感じの記事を探すなりしてね。

コードはここにおいた。

Haskell B-tree implementation · GitHub

まずはデータ型を定義。

data Tree a = Nil Int | Leaf Int [a] | Node Int [a] [Tree a] deriving Show

TreeNilLeafNodeであり、Nodeの中にはまたTreeのリストがあると再帰的に定義されている。Intの部分にはB-treeのオーダーが入る。Leafは保持するキーのリストを持つ。また、Nodeは中間ノードで、保持するキーのリストと保持する部分木のリストを持つ。

キーの検索

次にB-treeからキーを検索する関数を定義。すごく簡単だし直感的に書けるので Haskell かっこいい。

find :: (Ord a, Eq a) => Tree a -> a -> Bool
find (Nil _) _ = False
find (Leaf _ []) _ = False
find (Leaf m (k:ks)) x
  | x == k = True
  | x < k  = False
  | x > k  = find (Leaf m ks) x
find (Node _ [] (t:ts)) x = find t x
find (Node m (k:ks) (t:ts)) x
  | x == k = True
  | x < k  = find t x
  | x > k  = find (Node m ks ts) x

TreeNilである場合とLeafである場合とNodeである場合とで場合分けしている。Nilだった場合はキーは存在しないので単にFalseを返す。Leafだった場合は、保持しているキーのリストから再帰的に探す。Nodeだった場合はそのノードの中にキーがあればOK、なければ適切な部分木の中を再帰的に探索する。

ここでポイントなのがノード内におけるキーの探索方法。キーは昇順に格納されているので、左から順に探索して、xのほうが大きいうちはノード内のキーのリストと部分木のリストの残りの部分(ksts)に対して再帰的にfindを呼び。そしてxのほうがkより小さくなった時点でkに対応する部分木tに対して再帰的にfindを呼ぶ。感覚的に言えば、xのほうが大きい場合は横に移動して、xのほうが小さい場合は下に移動する。

キーの挿入

キーの挿入操作を実装する。まずはトップレベルの関数insertを定義する。

insert :: (Ord a, Eq a) => Tree a -> a -> Tree a
insert t x = if is_full t then insert_non_full (split t) x
                          else insert_non_full t x

まだ定義してない関数がいくつも出てきたけどこれらは順を追って説明する。insertではis_fullを使ってルートノードが full かどうか、つまりキーを (2m-1) 個持っているかどうか(m は B-tree のオーダー)で場合分けする。B-treeの定義より、各ノードは (m-1) 個以上、 (2m-1) 個以下のキーしか持てないのでこの場合分けが必要になる。

ルートノードが full だった場合はノードをsplitして full じゃない状態にしてからinsert_non_fullで実際に挿入する。もともとfull じゃなかった場合はそのままinsert_non_fullで挿入する。Introduction to Algorithms の実装方針はすこし特殊で、あるノードにキーを挿入する場合、挿入対象のノードをまず full じゃない状態にしてから挿入を行う。そうすることによって、ノードの分割処理が親方向に伝搬しないようになる。自分はこの実装方針のほうが理解しやすいと思った。

is_fullは単純で、引数として渡されたTreeが full かどうかを判定する。渡されたTreeNilなら常にFalseを返す。LeafNodeだった場合は保持しているキーの数(リストの長さ)が (2m-1) ならTrueを返す。

is_full :: (Ord a, Eq a) => Tree a -> Bool
is_full (Nil m) = False
is_full (Leaf m ks)
  | length ks == (2 * m - 1) = True
  | otherwise = False
is_full (Node m ks _)
  | length ks == (2 * m - 1) = True
  | otherwise = False

次にsplitを定義しよう。splitが何をしているのかは図を使って説明したほうが良さそうなので、 "Introduction to Algorithms" から図を引用して説明する*2

f:id:yamaguchiyuto:20180102225538p:plain

上図を見れば一発でわかると思う。つまり、split関数は full の状態のノードを3つのノードに分割してそれを返す。分割後の部分木のルートノードには分割前のノード内の真ん中のキーを格納する(上図では H)。そして、左の子ノードには分割前のノード内においてキー H より左のキーを、右の子ノードには分割前のノード内においてキー H より右のキーを格納する。

コードは次の通り。上で説明したことをそのまま実装しただけなので特にコードについて説明する点はないと思うけど、分割前のノードに格納されているキーのリストと部分木のリストを前半(first_half)と後半(last_half)に分割して、パターンマッチで新しい木を作り出して返してるだけ。

split :: (Ord a, Eq a) => Tree a -> Tree a
split (Leaf m keys) = Node m [k] [Leaf m k1, Leaf m k2]
  where k1 = first_half keys
        k:k2 = last_half keys
split (Node m keys trees) = Node m [k] [Node m k1 t1, Node m k2 t2]
  where k1 = first_half keys
        k:k2 = last_half keys
        t1 = first_half trees
        t2 = last_half trees

first_half :: [a] -> [a]
first_half xs = take (div (length xs) 2) xs

last_half :: [a] -> [a]
last_half xs = drop (div (length xs) 2) xs

最後に、メインとなるinsert_non_fullを定義しよう。名前の通り、この関数に引数として渡される木は full でないことが保証されている。

ちょっと長いけど、挿入先のノードがNilの場合、Leafの場合、Nodeの場合で場合分けしているだけなので、順を追って見てみよう。

insert_non_full :: (Ord a, Eq a) => Tree a -> a -> Tree a
insert_non_full (Nil m) x = Leaf m [x]
insert_non_full (Leaf m []) x = Leaf m [x]
insert_non_full l@(Leaf m keys@(k:ks)) x
  | x == k = l
  | x < k  = Leaf m (x:keys)
  | x > k  = Leaf m (k:newKs)
    where Leaf _ newKs = insert_non_full (Leaf m ks) x
insert_non_full (Node m [] (t:ts)) x = if is_full t then insert_non_full (split t) x
                                                    else Node m [] [(insert_non_full t x)]
insert_non_full n@(Node m keys@(k:ks) trees@(t:ts)) x
  | x == k = n
  | x < k  = if is_full t then insert_non_full (Node m (newK:k:ks) (newT1:newT2:ts)) x
                          else Node m keys ((insert_non_full t x):ts)
  | x > k  = Node m (k:newKs) (t:newTs)
    where Node _ newKs newTs = insert_non_full (Node m ks ts) x
          Node _ [newK] [newT1, newT2] = split t

まず、Nilに挿入する場合は単に挿入キーをひとつだけ持ったLeafを返す。これは単純。

次に、Leafに挿入する場合を見てみる。Leafに格納されているキーのリストが空の場合はそこに挿入キーを格納して新しいLeafを返す。キーのリストが空でない場合は、現在注目しているキーkと挿入キーxを比較する。x == kの場合はすでにキーが存在するので、何もせずにそのままlを返す。x < kの場合はxを挿入する場所が見つかったので、xkeysにつなげて返す。最後にx > kの場合はxを挿入する場所がまだ見つからないので、残りのキーksに対して再帰的にinsert_non_fullをする。そして帰ってきたLeafからパターンマッチでキーのリストnewKsを取り出して、kとくっつけて返す。

最後に、Nodeに挿入する場合について。B-treeではキーの挿入はかならずLeafに対して行うので、Nodeへの挿入は行わず "適切な子部分木" に対して再帰的にinsert_non_fullを呼ぶことになる。またinsert_non_fullを呼ぶ前には、対象となる部分木が full であるかどうかを必ずチェックして、 full であるなら split してから insert_non_fullをする。

"適切な子部分木" を見つける部分はLeafのときとほとんど同じなんだけど、x == kのときはすでにキーが存在するので何もせず、x < kのときは "適切な子部分木" が見つかったのでそこに再帰的にinsert_non_fullをし、x > kのときはまだ "適切な子部分木" が見つからないので残りのキーの中からの探索を続ける。

テスト

実装できたので簡単にテストする。オーダーm=2のB-treeを作る。まず適当に作ったキーのリストに対してfoldlで順にinsertして木を得る。pretty print はサボったので見づらいんだけど、ちゃんとB-treeの定義に従った深さ3の木ができているのが分かる。挿入したキーを探索すると全て見つかるし、挿入していないキーを探索するともちろん見つからない。うまく動いてるっぽい。

*Main> :l btree
[1 of 1] Compiling Main             ( btree.hs, interpreted )
Ok, one module loaded.
*Main> keys = [432,463,7,453,1,52,7,8,6,35,756,987,4454,765]
*Main> t = foldl insert (Nil 2) keys
*Main> t
Node 2 [432] [Node 2 [7] [Leaf 2 [1,6],Leaf 2 [8,35,52]],Node 2 [463,987] [Leaf 2 [453],Leaf 2 [756,765],Leaf 2 [4454]]]
*Main> map (find t) keys
[True,True,True,True,True,True,True,True,True,True,True,True,True,True]
*Main> find t 9999999
False

キーの削除

次回!

*1:ビッグデータ基盤勉強会

*2:p. 496より