107 lines
2.6 KiB
Lua
107 lines
2.6 KiB
Lua
---@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
|