ninewise/adventofcode-2021
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

generalize union find

Browse Source
This commit is contained in:
Felix Van der Jeugt 2021-12-09 11:30:02 +01:00
parent ef0775f5e2
commit e058039f35
No known key found for this signature in database
GPG Key ID: 58B209295023754D
2 changed files with 40 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
day09/part2.lua
View File

@ -10,6 +10,7 @@ for line in io.lines(arg[3]) do
end end
table.insert(hm, row) table.insert(hm, row)
end end
local w = #hm[1]
local function get(r, c) local function get(r, c)
return (hm[r] or {})[c] or 9 return (hm[r] or {})[c] or 9
@ -20,19 +21,22 @@ for r, row in pairs(hm) do
for c, col in pairs(row) do for c, col in pairs(row) do
if col ~= 9 then if col ~= 9 then
if get(r - 1, c) ~= 9 then if get(r - 1, c) ~= 9 then
union:makeset(r, c) union:makeset(r*w + c)
union:makeset(r - 1, c) union:makeset(r*w - w + c)
union:union(r, c, r - 1, c) union:union(r*w + c, r*w - w + c)
end end
if get(r, c - 1) ~= 9 then if get(r, c - 1) ~= 9 then
union:makeset(r, c) union:makeset(r*w + c)
union:makeset(r, c - 1) union:makeset(r*w + c - 1)
union:union(r, c, r, c - 1) union:union(r*w + c, r*w + c - 1)
end end
end end
end end
end end
local sizes = union:sizes() local sizes = {}
for set, size in union:unions() do
table.insert(sizes, size)
end
table.sort(sizes) table.sort(sizes)
print(sizes[#sizes] * sizes[#sizes - 1] * sizes[#sizes - 2]) print(sizes[#sizes] * sizes[#sizes - 1] * sizes[#sizes - 2])

57
unionfind.lua
View File

@ -1,42 +1,43 @@
local function makeset(self, r, c) local function makeset(self, x)
if self.forest[r] == nil then self.forest[r] = {} end if self.forest[x] == nil then
if self.forest[r][c] == nil then self.forest[x] = { x, 1 }
self.forest[r][c] = { r, c, 1 }
end end
end end
local function find(self, r, c) local function find(self, x)
if self.forest[r] == nil then self.forest[r] = {} end local x0, s0 = unpack(self.forest[x] or {})
local r0, c0, s0 = unpack(self.forest[r][c] or {}) if x == x0 then
if r0 == r and c0 == c then return x0, s0
return r0, c0, s0
else else
self.forest[r][c] = { self:find(r0, c0) } x0, s0 = self:find(x0)
return unpack(self.forest[r][c]) self.forest[x] = { x0, s0 }
return x0, s0
end end
end end
local function union(self, ar, ac, br, bc) local function union(self, x, y)
a0r, a0c, a0s = self:find(ar, ac) x0, xs = self:find(x)
b0r, b0c, b0s = self:find(br, bc) y0, ys = self:find(y)
if a0r == b0r and a0c == b0c then return end if x0 == y0 then return end
if a0s < b0s then if xs < ys then
a0r, a0c, a0s, b0r, b0c, b0s = b0r, b0c, b0s, a0r, a0c, a0s x0, y0 = y0, x0
end end
self.forest[b0r][b0c] = { a0r, a0c, a0s + b0s } self.forest[x0] = { x0, xs + ys }
self.forest[a0r][a0c] = { a0r, a0c, a0s + b0s } self.forest[y0] = { x0, xs + ys }
end end
local function sizes(self) function unions(self)
local sizes = {} local index, item
for r, row in pairs(self.forest) do return function()
for c, col in pairs(row) do repeat
if r == col[1] and c == col[2] then index, item = next(self.forest, index)
table.insert(sizes, col[3]) until item == nil or index == item[1]
if item ~= nil then
return unpack(item)
else
return nil
end end
end end
end
return sizes
end end
function unionfind() function unionfind()
@ -46,6 +47,6 @@ function unionfind()
makeset = makeset, makeset = makeset,
find = find, find = find,
union = union, union = union,
sizes = sizes, unions = unions,
} }
end end