mutable_priority_queue.rb

require 'minitest/autorun'

class Item
  include Comparable

  def initialize(item, priority:)
    @item = item
    @priority = priority
  end

  def ==(other)
    item == other.item && priority == other.priority
  end

  def <=>(other)
    priority <=> other.priority
  end

  protected

  attr_reader :item, :priority
end

class PriorityQueue
  def initialize
    @tree = []
  end

  def <<(n)
    @tree << n
    bubble_up last_index
    self
  end

  private def last_index
    @tree.size - 1
  end

  private def bubble_up(i)
    return if i.zero?
    parent = i / 2

    if bigger?(i, parent)
      swap i, parent
      bubble_up parent
    end
  end

  private def bigger?(left, right)
    return if left > last_index
    @tree[left] > @tree[right]
  end

  private def swap(left, right)
    @tree[left], @tree[right] = @tree[right], @tree[left]
  end

  def pop
    value = @tree[0]
    last_value = @tree.pop

    if @tree.size > 0
      @tree[0] = last_value
      bubble_down 0
    end

    value
  end

  private def bubble_down(i)
    left, right = i * 2, i * 2 + 1

    largest = i
    largest = left if bigger?(left, largest)
    largest = right if bigger?(right, largest)

    if largest != i
      swap i, largest
      bubble_down largest
    end
  end
end

class TestPriorityQueue < MiniTest::Test
  def test_1_always_pops_off_the_biggest_item
    heap = PriorityQueue.new
    heap << 11 << 5 << 8 << 3 << 15 << 4

    assert_equal 15, heap.pop
    assert_equal 11, heap.pop
    assert_equal 8, heap.pop
    assert_equal 5, heap.pop
    assert_equal 4, heap.pop
    assert_equal 3, heap.pop
    assert_equal nil, heap.pop
  end

  def test_2_always_pops_off_the_biggest_item
    heap = PriorityQueue.new
    heap << 3 << 5 << 1 << 7 << 6

    assert_equal 7, heap.pop
    assert_equal 6, heap.pop
    assert_equal 5, heap.pop
    assert_equal 3, heap.pop
    assert_equal 1, heap.pop
    assert_equal nil, heap.pop
  end

  def test_3_always_pops_off_the_biggest_item
    heap = PriorityQueue.new
    heap << 2 << 3 << 1 << 9 << 7 << 15 << 14

    assert_equal 15, heap.pop
    assert_equal 14, heap.pop
    assert_equal 9, heap.pop
    assert_equal 7, heap.pop
    assert_equal 3, heap.pop
    assert_equal 2, heap.pop
    assert_equal 1, heap.pop
    assert_equal nil, heap.pop
  end

  def test_pops_item_with_highest_priority
    queue = PriorityQueue.new

    queue << Item.new(1, priority: 2) <<
             Item.new(2, priority: 3) <<
             Item.new(3, priority: 1) <<
             Item.new(4, priority: 9) <<
             Item.new(5, priority: 7) <<
             Item.new(6, priority: 15) <<
             Item.new(7, priority: 14)

    assert_equal Item.new(6, priority: 15), queue.pop
    assert_equal Item.new(7, priority: 14), queue.pop
    assert_equal Item.new(4, priority: 9), queue.pop
    assert_equal Item.new(5, priority: 7), queue.pop
    assert_equal Item.new(2, priority: 3), queue.pop
    assert_equal Item.new(1, priority: 2), queue.pop
    assert_equal Item.new(3, priority: 1), queue.pop
  end
end