Rubyで平衡二分探索木を実装（AA木）

競プロで必要になったときのため ＆ アルゴリズムの勉強のため。

平衡二分探索木（self-balancing binary search tree）を調べたらいくつか出てきたが、基礎的な考え方をしていて且つ場合分けが少ない（と感じた）AA木を選ぶことにした。Wikipediaの記事にあった説明にほぼ沿って実装した。

※ ちなみに勉強前は、ふつうの二分探索木すらきちんと知らなかった。

実装

• Node クラス自身に、自分を木の頂点とした再帰的な処理を書いた。
  • 末端に特別な NULL_NODE オブジェクトを配置し、条件分岐を減らした。
  • 再帰を遡る必要の無い箇所には大域脱出を差し込んだ。（消しても動作する）
• 既に存在している値を挿入した際は、別のノードを作って記憶するようにした。
  • 以前の版の実装では同じ値ばかりのときに削除が遅かったため、削除ノードの検索に細工を加えた。
• コード末尾には簡単なテスト（デバッグおよび性能測定）を付けた。

aa_tree.rb

class AATree
	class Node
		attr_accessor :left, :right, :value, :level

		def initialize(value)
			@left = @right = NULL_NODE
			@value = value
			@level = 1
		end

		def search(x)
			if x < @value
				@left.search(x)
			elsif x > @value
				@right.search(x)
			else
				throw TAG_FOUND, self
				self
			end
		end

		def insert(x)
			if x > @value
				@right = @right.insert(x)
				throw TAG_BALANCED if @right.level == @level - 1
			else
				@left = @left.insert(x)
				throw TAG_BALANCED if @left.level == @level - 1
			end

			t = self
			t = t.skew
			t = t.split
			return t
		end

		def delete(x)
			if x > @value
				@right = @right.delete(x)
				throw TAG_BALANCED if @right.level == @level - 1
			elsif x < @value || x == @left.value
				@left = @left.delete(x)
				throw TAG_BALANCED if @left.level == @level - 1
			else
				return @right if @left == NULL_NODE  # may return NULL_NODE

				s = @right
				s = s.left while s.left != NULL_NODE  # successor node
				@value, s.value = s.value, @value
				@right = @right.delete(x)
				throw TAG_BALANCED if @right.level == @level - 1
			end

			# self.decrease_level
			limit = [@left.level, @right.level].min + 1
			@level = limit if @level > limit
			@right.level = limit if @right.level > limit

			t = self
			t = t.skew
			r = t.right = t.right.skew
			r.right = r.right.skew
			t = t.split
			t.right = t.right.split
			return t
		end

		def skew
			return self if @level > @left.level

			t = @left
			@left = t.right
			t.right = self
			return t
		end

		def split
			return self if @level > @right.right.level

			t = @right
			@right = t.left
			t.left = self
			t.level += 1
			return t
		end

		def inspect
			[@left, @value, @right].inspect
		end
	end

	class NullNode < Node
		def initialize
			@left = @right = self
			@value = nil
			@level = 0
		end

		def search(x)
			throw TAG_FOUND, self
			self
		end

		def insert(x)
			Node.new(x)
		end

		def delete(x)  # noop: tree does not have x
			throw TAG_BALANCED
			self
		end

		def skew; self; end
		def split; self; end

		def inspect; nil.inspect; end
	end

	NULL_NODE = NullNode.new
	TAG_FOUND = :found
	TAG_BALANCED = :balanced

	def initialize
		@root = NULL_NODE
	end

	def search(x)
		catch(TAG_FOUND) { @root.search(x) }.value
	end

	def insert(x)
		catch(TAG_BALANCED) { @root = @root.insert(x) }
		self
	end

	def delete(x)
		catch(TAG_BALANCED) { @root = @root.delete(x) }
		self
	end

	def empty?
		@root == NULL_NODE
	end
end

if $0 == __FILE__
	data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3]

	tree = AATree.new
	data.each { |x| p tree.insert(x) }
	p 10.times.map { |x| tree.search(x) }
	data.each { |x| p tree.delete(x) }
end

if $0 == __FILE__
	require 'securerandom'
	require 'benchmark'

	n = 10**5
	data = SecureRandom.random_bytes(n * 2).unpack("s*")  # int16_t
#	data = SecureRandom.random_bytes(n * 4).unpack("l*")  # int32_t
	tree = AATree.new

	Benchmark.bm do |bm|
		bm.report("insert") { data.each { |x| tree.insert(x) } }
		bm.report("search") { data.each { |x| tree.search(x) } }
		bm.report("delete") { data.each { |x| tree.delete(x) } }
	end

	raise Exception, "Wrong implementation" unless tree.empty?
end

勉強メモ

二分探索木全般の操作

あるノードを頂点とした部分木に対する操作として記述できる。再帰呼び出しして対象の部分木を小さくしていく考え方が楽。

• 探索
  • 二分探索木のデータ構造では、頂点のノードに格納された値より小さいものは左側の部分木、大きいものは右側の部分木にある。そのため探したい値と頂点ノードの値を比較して左右どちらに絞り込むか決めればいい。
  • 見ている部分木の頂点ノードの値が一致すれば発見となる。一方でnullノードまで到達したら、木にその値が無いということになる。
• 挿入
  • 探索と同じ要領で、適切な位置のnullノードを探して、値を持つノードに置き換える。
• 削除
  • 探索と同じ要領で、削除したい値を持つノードを探す。そのノードが葉（子ノードを持たない）ならそのまま削除していい。
  • 削除したいノードが葉でない（子ノードを持つ）なら、以下の方法で葉の削除に帰着させる。
    • 削除予定の値の「次」または「前」の値のノードを探す。例えば「次」は、右側の子ノードの一番左の子孫として発見できる。
    • 削除予定の値のノードと前後の値のノードとで、値を交換しておく。
    • 交換後の部分木から再び値を探して削除する。（これも葉でない可能性があるが、再帰で書けば対応できる）

平衡化の処理（一般論）

二分探索が効率よく働くのは、データ数が左右でほぼ均等に配置されている（木の高さが O(log n) に抑えられている）とき。挿入や削除の位置によっては木のバランスが悪くなっていき、以降の探索で効率が落ちてしまう。そのため、二分探索木を適宜変形してバランスを整える必要がある。この方法が色々と提案されていて、それぞれ処理効率などに特徴がある。

葉から根に向かってバランスを確認する際は、明示的に親ノードを参照して上るほかにも、前節の操作で再帰的に処理して戻る際に平衡化の処理を入れられる。探索から外れているノードは手を付けてなく平衡が保たれているはずなので、毎回全ての葉から出発して確認するのではなく、ノードにうまい指標を格納しておいて使うのが良い。また、根に向かう途中で平衡化が完了すればそこで打ち切っていい。