Create priority_queue.lua

This commit is contained in:
Ivan Yuriev 2025-09-07 14:27:02 +03:00
parent dd1d64506d
commit e8bb7306ac

View File

@ -0,0 +1,106 @@
---@class PriorityQueue
---@field private data any[] # внутренний массив-куча (индексация с 1)
---@field private cmp fun(a:any, b:any):boolean # компаратор: true, если a выше по приоритету, чем b
local PriorityQueue = {}
PriorityQueue.__index = PriorityQueue
---Создать очередь с приоритетом.
---@param cmp fun(a:any, b:any):boolean|nil # если nil, используется a < b (мин-куча)
---@return PriorityQueue
function PriorityQueue.new(cmp)
local self = setmetatable({}, PriorityQueue)
self.data = {}
self.cmp = cmp or function(a, b) return a < b end
return self
end
-- ===== Внутренние утилиты =====
---@param i integer @индекс узла
---@param j integer @индекс узла
function PriorityQueue:_swap(i, j)
self.data[i], self.data[j] = self.data[j], self.data[i]
end
---@param i integer
function PriorityQueue:_sift_up(i)
while i > 1 do
local p = math.floor(i / 2)
if self.cmp(self.data[i], self.data[p]) then
self:_swap(i, p)
i = p
else
break
end
end
end
---@param i integer
function PriorityQueue:_sift_down(i)
local n = #self.data
while true do
local l, r = i * 2, i * 2 + 1
local best = i
if l <= n and self.cmp(self.data[l], self.data[best]) then
best = l
end
if r <= n and self.cmp(self.data[r], self.data[best]) then
best = r
end
if best ~= i then
self:_swap(i, best)
i = best
else
break
end
end
end
-- ===== Публичный API =====
---Вставка элемента.
---@param x any
function PriorityQueue:insert(x)
table.insert(self.data, x)
self:_sift_up(#self.data)
end
---Посмотреть первый элемент без удаления.
---@return any|nil
function PriorityQueue:peek()
return self.data[1]
end
---Извлечь первый элемент (корень кучи).
---@return any|nil
function PriorityQueue:pop()
local n = #self.data
if n == 0 then return nil end
local top = self.data[1]
local last = self.data[n]
self.data[n] = nil
if n > 1 then
self.data[1] = last
self:_sift_down(1)
end
return top
end
---Текущее количество элементов.
---@return integer
function PriorityQueue:size()
return #self.data
end
---Пуста ли очередь.
---@return boolean
function PriorityQueue:is_empty()
return #self.data == 0
end
return PriorityQueue