Rubyでヒープの実装

2020年11月1日

Atcoderで問題解くのに優先度付きキューが必要ということで、実装しました。

最大ヒープを作成(不等号の向き変えれば最小にできます)
値の取り出しと挿入がどちらもlognなのは魅力的

class BinaryHeap
  def initialize(array=[])
    @heap = []
    array.each do |element|
      add(element)
    end
  end

  def add(value)
    @heap.push(value)
    up_heap
  end

  def extract_max
    max_value = @heap[0]
    @heap[0] = @heap.pop
    down_heap
    max_value
  end

  private
  def up_heap
    index = @heap.size - 1

    while index != 0
      if @heap[index] < @heap[parent_index(index)]
        break
      end
      swap(index, parent_index(index))
      index = parent_index(index)
    end
  end

  def down_heap
    index = 0

    while has_child?(index) do
      l = left_child_index(index)
      r = right_child_index(index)

      if @heap[r].nil?
        larger_index = l
      else
        larger_index = ( @heap[l] >= @heap[r]) ? l : r
      end

      if @heap[index] > @heap[larger_index]
        break
      end

      swap(index, larger_index)
      index = larger_index
    end

  end

  def parent_index(index)
    (index - (index.even? ? 2 : 1) )/2
  end

  def left_child_index(index)
    index * 2 + 1
  end

  def right_child_index(index)
    index * 2 + 2
  end

  def swap(i, j)
    @heap[i], @heap[j] = @heap[j], @heap[i]
  end

  def has_child?(index)
    ((index * 2) + 1) < @heap.size
  end
end

a = [3, 1, 4, 2 ,8 ,5, 6]

bh = BinaryHeap.new(a)

Wikiの図を見るとわかりやすい

言葉でかんたんに説明すると、
pop → 先頭を取り、末尾を先頭に設置。 そのあと比較しておろしていく
push → まず末尾につけ、親と比較して上までもっていく

実装したあとにgemがあることに気づきました・・・
https://github.com/general-CbIC/ruby-heap