class Heap

# File lib/rb_heap/sort.rb, line 17
def self.invertComparison(order, &compare_fn)
  if block_given?
    lambda{|a, b| not compare_fn.call(a, b)}
  elsif order == :< or order.nil?
    lambda{|a, b| a > b}
  elsif order == :>
    lambda{|a, b| a < b}
  else
    raise ArgumentError.new("The comparison symbol needs to be either :> or :<")
  end
end

# File lib/rb_heap/heap.rb, line 2
def initialize(compare_symbol = :<, storage = [], &compare_fn)
  @heap = storage
  @size = 0
  initialize_compare(compare_symbol, &compare_fn)
end

# File lib/rb_heap/sort.rb, line 2
def self.sort(array, order = :<, &compare_fn)
  compare_fn = self.invertComparison(order, &compare_fn)
  heap = Heap.new(order, array, &compare_fn)

  array.each do |element|
    heap.add(element)
  end

  (0...array.size).reverse_each do |i|
    array[i] = heap.pop
  end

  heap.to_a
end

# File lib/rb_heap/heap.rb, line 34
def add(element)
  @heap[@size] = element
  @size += 1
  rebalance_up(size - 1)
  self
end

# File lib/rb_heap/heap.rb, line 58
def clear
  @heap = []
  @size = 0
end

# File lib/rb_heap/heap.rb, line 10
def empty?
  size == 0
end

# File lib/rb_heap/heap.rb, line 48
def offer(element)
  if compare(peak, element)
    result = peak
    replace(element)
    result
  else
    element
  end
end

# File lib/rb_heap/heap.rb, line 14
def peak
  @heap[0]
end

# File lib/rb_heap/heap.rb, line 18
def pop
  result = peak

  if size > 1
    @size -= 1
    @heap[0] = @heap[@size]
    rebalance_down(0)
  else
    @size = 0
  end

  @heap[@size] = nil

  result
end

# File lib/rb_heap/heap.rb, line 43
def replace(element)
  @heap[0] = element
  rebalance_down(0)
end

# File lib/rb_heap/heap.rb, line 63
def to_a
  @heap.reject{|element| element.nil? }
end

# File lib/rb_heap/heap.rb, line 81
def compare(a, b)
  @compare.call(a, b)
end

# File lib/rb_heap/heap.rb, line 115
def has_left(i)
  left(i) < size
end

# File lib/rb_heap/heap.rb, line 107
def has_parent(i)
  i >= 1
end

# File lib/rb_heap/heap.rb, line 123
def has_right(i)
  right(i) < size
end

# File lib/rb_heap/heap.rb, line 69
def initialize_compare(symbol, &fn)
  @compare = if block_given?
    fn
  elsif symbol == :< or symbol.nil?
    lambda{|a, b| a < b}
  elsif symbol == :>
    lambda{|a, b| a > b}
  else
    raise ArgumentError.new("The comparison symbol needs to be either :> or :<")
  end
end

# File lib/rb_heap/heap.rb, line 119
def left(i)
  i * 2 + 1
end

# File lib/rb_heap/heap.rb, line 111
def parent(i)
  ((i - 1) / 2).floor
end

# File lib/rb_heap/heap.rb, line 94
def rebalance_down(i)
  left_i = left(i)
  right_i = right(i)

  if has_left(i) and compare(@heap[left_i], @heap[i]) and (not has_right(i) or compare(@heap[left_i], @heap[right_i]))
    @heap[i], @heap[left_i] = @heap[left_i], @heap[i]
    rebalance_down(left_i)
  elsif has_right(i) and compare(@heap[right_i], @heap[i])
    @heap[i], @heap[right_i] = @heap[right_i], @heap[i]
    rebalance_down(right_i)
  end
end

# File lib/rb_heap/heap.rb, line 85
def rebalance_up(i)
  parent_i = parent(i)

  if has_parent(i) and compare(@heap[i], @heap[parent_i])
    @heap[i], @heap[parent_i] = @heap[parent_i], @heap[i]
    rebalance_up(parent_i)
  end
end

# File lib/rb_heap/heap.rb, line 127
def right(i)
  i * 2 + 2
end