min_heap.tl

local record MinHeap<T>
   _size: integer
   _items: {T}
   _comparator: function(T, T):boolean

   size: integer

   metamethod __call: function(self, comparator?: function(T, T):boolean): MinHeap<T>
end

function MinHeap:__init(comparator: function(T, T):boolean)
   self._size = 0
   self._items = {}

   if comparator == nil then
      self._comparator = function(left:T, right:T):boolean
         -- cast to int just to suppress teal error
         return (left as integer) < (right as integer)
      end
   else
      self._comparator = comparator
   end
end

function MinHeap:_parent_index(i:integer):integer return math.floor(i/2) end
function MinHeap:_left_child_index(i:integer):integer return i*2 end
function MinHeap:_right_child_index(i:integer):integer return i*2 + 1 end

function MinHeap:has_parent(i:integer):boolean return self:_parent_index(i) >= 1 end
function MinHeap:has_left_child(i:integer):boolean return self:_left_child_index(i) <= self._size end
function MinHeap:has_right_child(i:integer):boolean return self:_right_child_index(i) <= self._size end

function MinHeap:parent(i:integer):T return self._items[self:_parent_index(i)] end
function MinHeap:left_child(i:integer):T return self._items[self:_left_child_index(i)] end
function MinHeap:right_child(i:integer):T return self._items[self:_right_child_index(i)] end

function MinHeap:_swap(i:integer, j:integer)
    self._items[i], self._items[j] = self._items[j], self._items[i]
end

function MinHeap:_heapify_up()
    local index = self._size
    while self:has_parent(index) and self._comparator(self._items[index], self:parent(index)) do
        self:_swap(self:_parent_index(index), index)
        index = self:_parent_index(index)
    end
end

function MinHeap:insert(key:T)
    self._size = self._size + 1
    self._items[self._size] = key
    self:_heapify_up()
end

function MinHeap:insert_all(keys:{T})
   for i=1,#keys do
      self:insert(keys[i])
   end
end

function MinHeap:_heapify_down()
    local index = 1
    while self:has_left_child(index) do
        local smallerChildIndex = self:_left_child_index(index)
        if self:has_right_child(index) and self._comparator(self:right_child(index), self:left_child(index)) then
            smallerChildIndex = self:_right_child_index(index)
        end

        if self._comparator(self._items[index], self._items[smallerChildIndex]) then
            break
        else
            self:_swap(index, smallerChildIndex)
        end
        index = smallerChildIndex
    end
end

function MinHeap:extract_min():T
    if self._size == 0 then error("Empty heap") end
    local item = self._items[1]
    self._items[1] = self._items[self._size]
    self._size = self._size - 1
    self:_heapify_down()
    return item
end

function MinHeap:peek():T
    if self._size == 0 then error("Empty heap") end
    return self._items[1]
end

return MinHeap