adventofcode-2021/unionfind.lua

53 lines
899 B
Lua
Raw Normal View History

2021-12-09 11:30:02 +01:00
local function makeset(self, x)
if self.forest[x] == nil then
self.forest[x] = { x, 1 }
2021-12-09 09:59:20 +01:00
end
end
2021-12-09 11:30:02 +01:00
local function find(self, x)
local x0, s0 = unpack(self.forest[x] or {})
if x == x0 then
return x0, s0
2021-12-09 09:59:20 +01:00
else
2021-12-09 11:30:02 +01:00
x0, s0 = self:find(x0)
self.forest[x] = { x0, s0 }
return x0, s0
2021-12-09 09:59:20 +01:00
end
end
2021-12-09 11:30:02 +01:00
local function union(self, x, y)
x0, xs = self:find(x)
y0, ys = self:find(y)
if x0 == y0 then return end
if xs < ys then
x0, y0 = y0, x0
2021-12-09 09:59:20 +01:00
end
2021-12-09 11:30:02 +01:00
self.forest[x0] = { x0, xs + ys }
self.forest[y0] = { x0, xs + ys }
2021-12-09 09:59:20 +01:00
end
2021-12-09 11:30:02 +01:00
function unions(self)
local index, item
return function()
repeat
index, item = next(self.forest, index)
until item == nil or index == item[1]
if item ~= nil then
return unpack(item)
else
return nil
2021-12-09 09:59:20 +01:00
end
end
end
function unionfind()
local forest = {}
return {
forest = forest,
makeset = makeset,
find = find,
union = union,
2021-12-09 11:30:02 +01:00
unions = unions,
2021-12-09 09:59:20 +01:00
}
end