Priority Queueの練習

priorityqueue.rb

# Priority Queueの練習
#
# 参考
# - https://ufcpp.net/study/algorithm/col_heap.html

class PriorityQueue
  attr_reader :data, :test_function

  def initialize(test_function = ->(a, b) { a <=> b })
    @data = []
    @test_function = test_function
  end

  def push(value)
    insert(value)
  end
  alias_method  :<<, :push

  def append(array)
    array.each { |value| insert(value) }
  end

  def shift
    ret = @data[0]
    @data[0] = @data.last
    @data.pop
    make_heap
    ret
  end

  def first
    @data[0]
  end

  def size
    @data.size
  end

private
  def insert(value)
    @data << value

    i = @data.size - 1
    while (i > 0)
      parent = ((i - 1) / 2).floor
      break unless test_function.call(@data[parent], @data[i]) < 0
      @data[parent], @data[i] = @data[i], @data[parent]
      i = parent
    end
  end

  def make_heap
    parent = 0
    loop {
      child = 2 * parent + 1
      break unless child < @data.size

      child += 1 if child < @data.size - 1 and (test_function.call(@data[child], @data[child + 1]) < 0)

      if test_function.call(@data[parent], @data[child]) < 0
        @data[parent], @data[child] = @data[child], @data[parent]
      end
      parent = child
    }
  end
end

queue = PriorityQueue.new
queue << 3
queue << 1
queue << 4
queue.append([1, 5, 9, 2])

puts "Max: #{queue.first}" 
pp Array.new(queue.size) { queue.shift }