Ruby で Priority Queue を実装してみたい

厳密にはやらないでラフにやります。

• 集合に対して要素を追加できる
• 集合から最小の要素を取り出せる

この二つの機能を持った class を安直に書いてみる。

# 安直につくったもの（パフォーマンス悪い）
class PriorityQueue
  def initialize(array = [])
    @data = []
    array.each{|a| push(a)}
  end

  def push(element)
    @data.push(element)
  end

  def pop
    @data.sort!.shift
  end
end

こんな感じで動きます。

pq = PriorityQueue.new([3,5,1])
pq.push(4)
pq.pop # => 1

heap というデータ構造をつかって書いてみる。
以下、雑なので時間があったら丁寧にしたい。（丁寧版は こちらの gist に書くことにした）

# heap を使ったもの（パフォーマンス良い）
class PriorityQueue
  def initialize(array = [])
    @data = []
    array.each{|a| push(a)}
  end

  def push(element)
    @data.push(element)
    bottom_up
  end

  def pop
    if size == 0
      return nil
    elsif size == 1
      return @data.pop
    else
      min = @data[0]
      @data[0] = @data.pop
      top_down
      return min
    end
  end

  private

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

  def size
    @data.size
  end

  def parent_idx(target_idx)
    (target_idx - (target_idx.even? ? 2 : 1)) / 2
  end

  def bottom_up
    target_idx = size - 1
    return if target_idx == 0
    parent_idx = parent_idx(target_idx)
    while (@data[parent_idx] > @data[target_idx])
      swap(parent_idx, target_idx)
      target_idx = parent_idx
      break if target_idx == 0
      parent_idx = parent_idx(target_idx)
    end
  end

  def top_down
    target_idx = 0

    # child がある場合
    while (has_child?(target_idx))

      a = left_child_idx(target_idx)
      b = right_child_idx(target_idx)

      if @data[b].nil?
        c = a
      else
        c = @data[a] <= @data[b] ? a : b
      end

      if @data[target_idx] > @data[c]
        swap(target_idx, c)
        target_idx = c
      else
        return
      end
    end
  end

  # @param Integer
  # @return Integer
  def left_child_idx(idx)
    (idx * 2) + 1
  end

  # @param Integer
  # @return Integer
  def right_child_idx(idx)
    (idx * 2) + 2
  end

  # @param Integer
  # @return Boolent
  def has_child?(idx)
    ((idx * 2) + 1) < @data.size
  end
end

ベンチマークをとってみる

# ベンチマークをとってみる
require 'benchmark'

@test_data = 100_000.times.map{rand(100_000)}
@test_cnt = 1_000

Benchmark.bm 10 do |r|
  r.report "Normal" do
    pq = PriorityQueue.new(@test_data)
    @test_cnt.times do
      pq.push(rand(100_000))
      pq.pop
    end
  end

  r.report "Heap" do
    pq = PriorityQueueWithHeap.new(@test_data)
    @test_cnt.times do
      pq.push(rand(100_000))
      pq.pop
    end
  end
end

# 結果
                 user     system      total        real
Normal       5.460000   0.010000   5.470000 (  5.494386)
Heap         0.070000   0.000000   0.070000 (  0.073038)

無事早くなって嬉しい。
コードは雑なので、時間があるときに丁寧なものにしたい。

ヒープの構造は蟻本を読んで実装しました。