RubyでPriority Queue

精進でこの問題を解いてたらPriority Queueがでてきました。

atcoder.jp

"Ruby Priority Queue"でググれば実装はたくさんでてくるのですが、実際自分で組んだ方が身につくと思い実装してみました。

github.com

MaxHeapかMinHeapかを外部から注入するのについsendを使いがち。 だいたいこんなふうに動きます。

 RSpec.describe 'PriorityQueue' do
  context 'direction is max' do
    it 'enqueue element' do
      pq = PriorityQueue.new(:max, [])

      pq.enqueue(3)
      pq.enqueue(9)
      pq.enqueue(6)
      pq.enqueue(1)

      expect(pq.heap.nodes).to eq([9, 3, 6, 1])
    end

    it 'dequeue element' do
      pq = PriorityQueue.new(:max, [9, 3, 6, 1])

      expect(pq.dequeue).to eq(9)
      expect(pq.heap.nodes).to eq([6, 3, 1])
    end
  end

  context 'direction is min' do
    it 'enqueue element' do
      pq = PriorityQueue.new(:min, [])

      pq.enqueue(3)
      pq.enqueue(9)
      pq.enqueue(6)
      pq.enqueue(1)

      expect(pq.heap.nodes).to eq([1, 3, 6, 9])
    end

    it 'dequeue element' do
      pq = PriorityQueue.new(:min, [1, 3, 6, 9])

      expect(pq.dequeue).to eq(1)
      expect(pq.heap.nodes).to eq([3, 9, 6])
    end
  end
end
 

Priority Queueというかヒープの特徴は

  • 最小値、最大値を探す時に全探査だとO(n)のところがO(1)で取れる
  • 追加、再構築はヒープの高さ(この場合二分ヒープ)がlog 2 nなので計算時間はO(log n)
  • 全探査でいちいち最大値、最小値を探すよりもO(n) - O( 1 + 2 (log n))分早い?
  • 何度も数字をいじってそのたびに最小値、最大値を求めるときに便利。

書いててわかってるか不安になってきた。