勉強のために Ruby で二分探索木を書いてみました。
二分探索木の Ruby 実装 · GitHub
なお、以下のコードは説明のためなので、Gist のコードを多少簡略化しています。
木(Tree)のクラスと要素(Node)のクラスを定義します。Node は Struct クラスを使って簡単に定義します。key と value の値が使えます。Node を追加するたび、key の順序で整序します。
Node = Struct.new(:key, :value, :left, :right)
Tree クラスは、要素の追加と探索はこんな風になります。
class Tree
def initialize
@root = nil
end
def insert(key, value)
unless @root
@root = Node.new(key, value)
return
end
node = @root
while node
if key < node.key
if node.left
node = node.left
else
node.left = Node.new(key, value)
break
end
elsif key > node.key
if node.right
node = node.right
else
node.right = Node.new(key, value)
break
end
else
if block_given?
node.value = yield(key, node.value, value)
else
node.value = value
end
break
end
end
end
def []=(key, value)
insert(key, value)
end
def search(key)
node = @root
while node
if key < node.key
node = node.left
elsif key > node.key
node = node.right
else
return node
end
end
nil
end
def [](key)
search(key)&.value
end
これだけでこんな風に使えます。まず、t[key] = value
という具合に、Node を登録しています。
$ pry
[1] pry(main)> require './binary_search_tree'
=> true
[2] pry(main)> t = Tree.new
=>
[3] pry(main)> t[5] = :いちご
=> :いちご
[4] pry(main)> t[3] = :みかん
=> :みかん
[5] pry(main)> t[10] = :りんご
=> :りんご
[6] pry(main)> t
=>
t[key]
でハッシュのように value を見ることができます。
[7] pry(main)> t[10]
=> :りんご
[8] pry(main)> t[5]
=> :いちご
メソッド print で木を簡単に出力します。
[9] pry(main)> t.print
-- 5 -- 3
-- 10
さらに Node を追加します。key の順序で自動的に整序されているのがわかると思います。
[10] pry(main)> t[4] = :柿
=> :柿
[11] pry(main)> t.print
-- 5 -- 3 --
-- 4
-- 10
メソッド each で木をトラバースします。
[12] pry(main)> t.each {|k, v| puts "#{k} --> #{v}"}
3 --> みかん
4 --> 柿
5 --> いちご
10 --> りんご
なお、each メソッドが定義されているので、Enumerable を include してあります。Enumerable のすべてのメソッドが使えます。
むずかしいのは Node の削除です。二分探索木の特性を壊さないように削除しなければなりません。コードはこんな感じ。
class Tree
def delete(key)
delete_min = ->(node) {
return node.right unless node.left
node.left = delete_min.(node.left)
node
}
del = ->(node) {
value = nil
if node
if key == node.key
value = node.value
if node.left.nil?
return node.right, value
elsif node.right.nil?
return node.left, value
else
min_node = search_min(node.right)
node.key = min_node.key
node.value = min_node.value
node.right = delete_min.(node.right)
end
elsif key < node.key
node.left , value = del.(node.left)
else
node.right, value = del.(node.right)
end
end
return node, value
}
@root, value = del.(@root)
value
end
end
ここで呼ばれている search_min メソッドは、それより下のすべての Node の中から、key が最小のものを探してくるものです。
使ってみます。root を削除してみます。
[13] pry(main)> t.delete(5)
=> :いちご
[14] pry(main)> t.print
-- 10 -- 3 --
-- 4
--
ちゃんと規則を壊さぬように削除できています。
最大、最小の key を探索できます。二分探索なので、一般的に高速になります。
[15] pry(main)> t.minimum
=> [3, :みかん]
他にもまだいろいろなメソッドが定義してありますが、これくらいが最小限度の機能だと思います。
メソッド一覧
- insert(key, value)
- []= はこれを使って実装してあるが、insert の場合はブロックを取って、キーの重複があった場合の処理ができる。
- []=(key, value)
- insert とほぼ同じだが、key が重複した場合は上書きする。
- search(key)
- [](key)
- delete(key)
- preorder(key = nil, &bk)
- 行きがけ。key を省略した場合は root からトラバースする。
- inorder(key = nil, &bk)
- postorder(key = nil, &bk)
- each(&bk)
- inorder を使って実装してある。不都合なら再定義する。
- breadth_first_search(key = nil, &bk)
- bfs(key = nil, &bk)
- breadth_first_search のエイリアス。
- minimum
- maximum
- inspect
- right_rotate(key)
- left_rotate(key)
- parent_by_node(target)
- parent(key)
- add_hash(hash)
- print