優先度付きキューにも使われる二分ヒープ構造をRubyで実装してみる

アルゴリズム図鑑 絵で見てわかる26のアルゴリズム

アルゴリズム図鑑を眺めていて、二分ヒープ構造は優先度付きキューに使われることを知った。面白いなーと思うと同時に、そういえば二分ヒープ構造の実装をしたことがなく、あまり理解できていないことに気づいた。そこで簡単に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すると必ず値が低い順に取得することができる。例えば値に優先度を示す数字を入れておけば、優先度付きキューとして使える。便利。