modlib/hashheap.lua
2022-01-19 17:17:07 +01:00

118 lines
3.0 KiB
Lua

-- Localize globals
local assert, math_floor, setmetatable, table_insert = assert, math.floor, setmetatable, table.insert
-- Set environment
-- Min. heap + Lua hash table to allow updating the stored values
local _ENV = {}
setfenv(1, _ENV)
local metatable = { __index = _ENV }
function less_than(a, b)
return a < b
end
--> empty, duplicate-free min heap with priority queue functionality
function new(less_than)
return setmetatable({ less_than = less_than, indices = {} }, metatable)
end
local function swap(self, child_index, parent_index)
local child_value, parent_value = self[child_index], self[parent_index]
self.indices[parent_value], self.indices[child_value] = child_index, parent_index
self[parent_index], self[child_index] = child_value, parent_value
end
local function heapify_up(self, index)
if index == 1 then
return
end
local parent_index = math_floor(index / 2)
if self.less_than(self[index], self[parent_index]) then
swap(self, index, parent_index)
heapify_up(self, parent_index)
end
end
local function heapify_down(self, index)
local left_child = index * 2
if left_child > #self then
return
end
local smallest_child = left_child + 1
if smallest_child > #self 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
swap(self, smallest_child, index)
heapify_down(self, smallest_child)
end
end
function push(self, value)
table_insert(self, value)
local last = #self
self.indices[value] = last
heapify_up(self, last)
end
function top(self)
return self[1]
end
-- TODO what if empty?
function pop(self)
local value = self[1]
self.indices[value] = nil
local last = #self
if last == 1 then
self[1] = nil
return value
end
self[1], self[last] = self[last], nil
heapify_down(self, 1)
return value
end
function find_index(self, element)
return self.indices[element]
end
-- Notify heap that the element has been decreased
function decrease(self, element)
heapify_up(self, assert(self:find_index(element)))
end
-- Notify heap that the element has been increased
function increase(self, element)
heapify_down(self, assert(self:find_index(element)))
end
-- Replaces the specified element - by identity - with the new element
function replace(self, element, new_element)
local index = assert(self:find_index(element))
assert(self:find_index(new_element) == nil, "no duplicates allowed")
self[index] = new_element
self.indices[element] = nil
self.indices[new_element] = index;
(self.less_than(new_element, element) and heapify_up or heapify_down)(self, index)
end
function remove(self, element)
local index = assert(self:find_index(element), "element not found")
self.indices[element] = nil
if index == #self then
self[index] = nil
else
local last_index = #self
local last_element = self[last_index]
self[last_index] = nil
self[index] = last_element
self.indices[last_element] = index;
(self.less_than(last_element, element) and heapify_up or heapify_down)(self, index)
end
end
-- Export environment
return _ENV