1
0
mirror of https://github.com/Refactorio/RedMew.git synced 2024-12-14 10:13:13 +02:00
RedMew/utils/priority_queue.lua
2019-06-22 22:53:58 +01:00

105 lines
2.7 KiB
Lua

local Debug = require 'utils.debug'
local is_closure = Debug.is_closure
local floor = math.floor
local PriorityQueue = {}
local function default_comparator(a, b)
return a < b
end
--- Min heap implementation of a priority queue. Smaller elements, as determined by the comparator,
-- have a higher priority.
-- @param comparator <function|nil> the comparator function used to compare elements
-- @usage
-- local PriorityQueue = require 'utils.priority_queue'
--
-- local queue = PriorityQueue.new()
-- PriorityQueue.push(queue, 4)
-- PriorityQueue.push(queue, 7)
-- PriorityQueue.push(queue, 2)
--
-- game.print(PriorityQueue.pop(queue)) -- 2
-- game.print(PriorityQueue.pop(queue)) -- 4
-- game.print(PriorityQueue.pop(queue)) -- 7
function PriorityQueue.new(comparator)
if comparator == nil then
comparator = default_comparator
elseif is_closure(comparator) then
error('comparator cannot be a closure.', 2)
end
return {_comparator = comparator}
end
local function heapify_from_end_to_start(self)
local comparator = self._comparator
local pos = #self
while pos > 1 do
local parent = floor(pos * 0.5)
local a, b = self[pos], self[parent]
if comparator(a, b) then
self[pos], self[parent] = b, a
pos = parent
else
break
end
end
end
local function heapify_from_start_to_end(self)
local comparator = self._comparator
local parent = 1
local smallest = 1
local count = #self
while true do
local child = parent * 2
if child > count then
break
end
if comparator(self[child], self[parent]) then
smallest = child
end
child = child + 1
if child <= count and comparator(self[child], self[smallest]) then
smallest = child
end
if parent ~= smallest then
self[parent], self[smallest] = self[smallest], self[parent]
parent = smallest
else
break
end
end
end
--- Returns the number of the number of elements in the priority queue.
function PriorityQueue.size(self)
return #self
end
-- Inserts an element into the priority queue.
function PriorityQueue.push(self, element)
self[#self + 1] = element
heapify_from_end_to_start(self)
end
-- Removes and returns the highest priority element from the priority queue.
function PriorityQueue.pop(self)
local element = self[1]
self[1] = self[#self]
self[#self] = nil
heapify_from_start_to_end(self)
return element
end
-- Returns, without removing, the highest priority element from the priority queue.
function PriorityQueue.peek(self)
return self[1]
end
return PriorityQueue