MaxHeap.lua

MaxHeap = {}

MaxHeap.data = {}

-- Convenience function
function table.swap(t, a, b)
    t[a], t[b] = t[b], t[a]
end

-- Convenience function to explode some text
local function explode(str, separator)
    local out = {}
    
    for w in str:gmatch("([^" .. separator .. "]+)") do out[#out + 1] = w end
    
    return out
end

--- Construct a new MaxHeap object
-- @return A new MaxHeap object
function MaxHeap:new()
    local obj = {}
    setmetatable(obj, self)
    self.__index = self
    
    return obj
end

--- Builds a MaxHeap from a table of values
-- @param list The table to use
function MaxHeap:from(list)
    for _, v in ipairs(list) do
        self:insert(v)
    end
end

--- "Bubbles" a node upward, recursively
-- @param index The index to start moving up
function MaxHeap:siftUp(index)
    local current = self.data[index]
    local parent, parentIndex = self:getParent(index)

    if (parent and parent < current) then
        table.swap(self.data, index, parentIndex)

        self:siftUp(parentIndex)
    end
end

--- Inserts a value into the heap using Williams' method
-- @param value The value to insert
function MaxHeap:insert(value)
    local i = self:size() + 1
    self.data[i] = value

    local parent, parentIndex = self:getParent(i)

    self:siftUp(i)
end

--- Get the size of the data set
-- @return Number of elements in the heap
function MaxHeap:size()
    return #self.data
end

function MaxHeap:getLeft(index)
    local newIndex = 2 * index

    return self.data[newIndex], newIndex
end

function MaxHeap:getRight(index)
    local newIndex = 2 * index + 1

    return self.data[newIndex], newIndex
end

function MaxHeap:getParent(index)
    local newIndex = math.floor(index / 2)

    return self.data[newIndex], newIndex
end

function MaxHeap:maxLookup()
    return self.data[1]
end

--- Pop the max value off the heap and restructure
-- @return The max value
function MaxHeap:extractMax()
    local max = self.data[1]
    
    self.data[1] = self.data[#self.data]
    
    self.data[#self.data] = nil
    
    self:fix()
    
    return max
end

--- Remove an arbitrary index from the heap
-- @param index The index to remove
-- @return The removed value, in case you wanted it
function MaxHeap:remove(index)
    if (index == 1) then
        return self:extractMax()
    end

    local o = self.data[index]
    local last = self.data[self:size()]
    
    self.data[index] = last
    self.data[self:size()] = nil
    
    local parent, parentIndex = self:getParent(index)
    
    if (last < parent) then
        self:siftUp(index)
    else
        self:heapify(index)
    end
    
    return o
end

--- Let there be heaps. Heapify a specific node and its subtrees.
-- @param index The index to start at
function MaxHeap:heapify(index)
    index = math.floor(index)
    local left, leftIndex = self:getLeft(index)
    local right, rightIndex = self:getRight(index)
    local size = self:size()

    if (index >= size / 2 and index <= size) then
        return
    end

    if (self.data[index] < left or self.data[index] < right) then 
        local temp = self.data[index]

        if (right > left) then
            table.swap(self.data, index, rightIndex)

            self:heapify(rightIndex)
        else
            table.swap(self.data, index, leftIndex)

            self:heapify(leftIndex)
        end
    end
end

-- If the data set is a binary tree, use Floyd's implementation to convert it
function MaxHeap:build()
    for i = self:size() / 2, 1, -1 do
        self:heapify(i)
    end
end

-- Make it purty again
function MaxHeap:fix()
    self:heapify(1)
end

-- local heap = MaxHeap:new()

-- local text = nil

-- local commands = {
-- 	size		= function(args) io.stdout:write(heap:size() .. "\n")		end,
-- 	insert		= function(args) heap:insert(tonumber(args[2]))				end,
-- 	delete		= function(args) heap:remove(tonumber(args[2]))				end,
-- 	extractmax	= function(args) heap:extractMax()							end,
-- 	maxlookup	= function(args) io.stdout:write(heap:maxLookup() .. "\n")	end			
-- }

-- while (text ~= "") do
--     local text = io.stdin:read()

--     if (not text or text == "") then break end
--  
--     local packed = explode(text, " ")
--     local command = string.lower(packed[1])

--     commands[command](packed)
-- end