Binary Heap

binary-heap.lua

function BinaryHeap(cmp)
    local heap = {}

    local function push(el)
        table.insert(heap, el)
        local i = #heap
        while i > 1 do
            local parent = math.floor(i / 2)
            if cmp(heap[parent], heap[i]) then
                break
            end
            heap[i], heap[parent], i = heap[parent], heap[i], parent
        end
    end

    local function pop()
        if #heap == 0 then
            return
        end
        local size = #heap
        local top = heap[1]
        heap[1], heap[size] = heap[size], heap[1]
        table.remove(heap, size)
        local i = 1
        while i < size-1 do
            local left = i * 2
            local right = i * 2 + 1
            local child
            if right <= size and cmp(heap[right], heap[left]) then
                child = right
            elseif left <= size then
                child = left
            else
                break
            end
            if cmp(heap[i], heap[child]) then
                break
            end
            heap[i], heap[child], i = heap[child], heap[i], child
        end
        return top
    end

    return {
        size = function () return #heap end,
        push = push,
        pop = pop
    }
end