adventofcode-2021/pqueue.lua

50 lines
1.2 KiB
Lua
Raw Permalink Normal View History

2021-12-15 10:51:53 +01:00
local function add(self, item)
table.insert(self.items, item)
local i, p = #self.items
local p = math.floor(i / 2)
while i > 1 and self.smaller(self.items[i], self.items[p]) do
self.items[p], self.items[i] = self.items[i], self.items[p]
i, p = p, math.floor(p / 2)
end
end
local function pop(self)
local item = self.items[1]
self.items[1] = self.items[#self.items]
table.remove(self.items)
local i = 1
while i * 2 <= #self.items do
if i * 2 + 1 <= #self.items then
if self.smaller(self.items[i*2], self.items[i])
or self.smaller(self.items[i*2+1], self.items[i]) then
if self.smaller(self.items[i*2], self.items[i*2+1]) then
self.items[i], self.items[i*2], i
= self.items[i * 2], self.items[i], i * 2
else
self.items[i], self.items[i*2+1], i
= self.items[i*2+1], self.items[i], i * 2 + 1
end
else
break
end
else
if self.smaller(self.items[i*2], self.items[i]) then
self.items[i], self.items[i * 2]
= self.items[i * 2], self.items[i]
end
i = i * 2
end
end
return item
end
function pqueue(smaller)
return {
items = {},
smaller = smaller
or function(a, b) return a < b end,
add = add,
pop = pop,
}
end