min tree aka priority queue in ruby for leetcode

min_tree.rb

class MinHeap
  def initialize(elements = [])
    @heap = []
    elements.each(&method(:push))
  end

  # `push` adds an element to the Min heap
  def push(element)
    heap << element
    rebalance_push(heap.size - 1)
  end

  # `pop` removes smallest element from the Min heap
  def pop
    heap[0], heap[heap.size - 1] = heap[heap.size - 1], heap[0]
    el = heap.pop
    rebalance_pop(0)
    el
  end

  # `peek` returns first(minimum) element from the Min heap
  def peek
    heap.first
  end

  # `pp` pretty prints a heap into console
  def pp(level: 0, index: 0)
    return unless heap[index] # poor man out of bound check

    n_el = 2**level
    puts("[#{level}] #{heap[index...index + n_el]}")
    pp(level: level + 1, index: index + n_el)
  end

  private

  attr_reader :heap

  def rebalance_push(index)
    parent_index = (index - 1) / 2
    return if index <= 0 || heap[parent_index] <= heap[index]

    heap[index], heap[parent_index] = heap[parent_index], heap[index]
    rebalance_push(parent_index)
  end

  def rebalance_pop(index)
    left_index = index * 2 + 1
    right_index = index * 2 + 2
    smallest = index

    smallest = left_index  if heap[left_index] && heap[left_index] < heap[smallest]
    smallest = right_index if heap[right_index] && heap[right_index] < heap[smallest]
    return if index == smallest # we already good

    heap[index], heap[smallest] = heap[smallest], heap[index]
    rebalance_pop(smallest) # down we go...
  end
end