アルゴリズム図鑑を眺めていて、二分ヒープ構造は優先度付きキューに使われることを知った。面白いなーと思うと同時に、そういえば二分ヒープ構造の実装をしたことがなく、あまり理解できていないことに気づいた。そこで簡単にRubyで実装をしてみたのでメモ。簡単なテストケースを作ったので多分合ってると思うけど、もしかしたらバグっているかも...
二分ヒープの詳細は二分ヒープ - Wikipediaも参考。
【2023/01/03 14:01追記】要素数が1の時に要素が空にならないバグがあったので修正しました。コメントありがとうございます。https://github.com/shibayu36/algorithms/commit/6c2ce588f7bc7fb890c6a560c7ab062c6f531a9a
実装
class Heap
attr_reader :nodes
def initialize(arr)
@nodes = []
arr.each { |val| push(val) }
end
def size
@nodes.size
end
def push(val)
@nodes.push(val)
current_index = size - 1
parent_index = parent_of(size - 1)
while current_index > 0 && @nodes[parent_index] > @nodes[current_index]
swap(parent_index, current_index)
current_index = parent_index
parent_index = parent_of(current_index)
end
end
def pop
return @nodes.pop if size <= 1
res = @nodes[0]
@nodes[0] = @nodes.pop
current_index = 0
child_index = min_child_of(current_index)
while child_index && @nodes[child_index] < @nodes[current_index]
swap(child_index, current_index)
current_index = child_index
child_index = min_child_of(current_index)
end
res
end
private
def parent_of(index)
(index - 1) / 2
end
def min_child_of(index)
left_child_index = (2 * index) + 1
right_child_index = (2 * index) + 2
left_val = nodes[left_child_index]
right_val = nodes[right_child_index]
if left_val && right_val
left_val <= right_val ? left_child_index : right_child_index
elsif left_val
left_child_index
end
end
def swap(i, j)
@nodes[i], @nodes[j] = @nodes[j], @nodes[i]
end
end
テスト
require 'heap'
describe Heap do
describe '#push' do
it do
heap = Heap.new([4, 1, 3])
expect(heap.nodes).to eq([1, 4, 3])
heap.push(5)
expect(heap.nodes).to eq([1, 4, 3, 5])
heap.push(0)
expect(heap.nodes).to eq([0, 1, 3, 5, 4])
end
end
describe '#pop' do
it do
heap = Heap.new([0, 1, 3, 5, 4, 3])
expect(heap.pop).to eq 0
expect(heap.nodes).to eq([1, 3, 3, 5, 4])
expect(heap.pop).to eq 1
expect(heap.nodes).to eq([3, 4, 3, 5])
expect(heap.pop).to eq 3
expect(heap.pop).to eq 3
expect(heap.pop).to eq 4
expect(heap.pop).to eq 5
expect(heap.pop).to eq nil
end
end
end
これはminヒープという構造。popすると必ず値が低い順に取得することができる。例えば値に優先度を示す数字を入れておけば、優先度付きキューとして使える。便利。