modlib/heap.lua
2021-03-27 20:10:49 +01:00

50 lines
1.1 KiB
Lua

local metatable = {__index = getfenv(1)}
function less_than(a, b) return a < b end
--> empty min heap
function new(less_than)
return setmetatable({less_than = less_than}, metatable)
end
function push(self, value)
table.insert(self, value)
local function heapify(index)
if index == 1 then
return
end
local parent = math.floor(index / 2)
if self.less_than(self[index], self[parent]) then
self[parent], self[index] = self[index], self[parent]
heapify(parent)
end
end
heapify(#self)
end
function pop(self)
local value = self[1]
local last = #self
if last == 1 then
self[1] = nil
return value
end
self[1], self[last] = self[last], nil
last = last - 1
local function heapify(index)
local left_child = index * 2
if left_child > last then
return
end
local smallest_child = left_child + 1
if smallest_child > last or self.less_than(self[left_child], self[smallest_child]) then
smallest_child = left_child
end
if self.less_than(self[smallest_child], self[index]) then
self[index], self[smallest_child] = self[smallest_child], self[index]
heapify(smallest_child)
end
end
heapify(1)
return value
end