70 lines
2.1 KiB
Lua
70 lines
2.1 KiB
Lua
#!/usr/bin/env lua
|
|
--[[
|
|
Representation in both ways:
|
|
The input mentions what bags a bag can contain,
|
|
we store both what bags a bag can contain and
|
|
what bags can contain a bag and in which amount.
|
|
Part 2 requires the former and part 1 requires the latter.
|
|
--]]
|
|
local containers = {}
|
|
local contains = {}
|
|
|
|
for line in io.lines() do
|
|
local bag_end, contain_start = line:find(" bags contain ")
|
|
local bag_color = line:sub(1, bag_end - 1)
|
|
contains[bag_color] = {}
|
|
|
|
for count, color in line:sub(contain_start + 1):gmatch("(%d+) (%w+ %w+) bag") do
|
|
if not containers[color] then
|
|
containers[color] = {}
|
|
end
|
|
containers[color][bag_color] = count
|
|
contains[bag_color][color] = count
|
|
end
|
|
end
|
|
|
|
--[[
|
|
For part 1, we build a set of bags that can contain the shiny gold bag.
|
|
We will iterate over and over on this set until we no longer find any new
|
|
bag colors that can contain others bags that can contain the shiny gold bag.
|
|
--]]
|
|
local part1set, found = {}, 0
|
|
for color, _ in pairs(containers["shiny gold"] or {}) do
|
|
part1set[color] = true
|
|
found = found + 1
|
|
end
|
|
|
|
local found_in_iteration = 0
|
|
repeat
|
|
found_in_iteration = 0
|
|
for known_color, _ in pairs(part1set) do
|
|
for color, _ in pairs(containers[known_color] or {}) do
|
|
if not part1set[color] then
|
|
found_in_iteration = found_in_iteration + 1
|
|
found = found + 1
|
|
part1set[color] = true
|
|
end
|
|
end
|
|
end
|
|
until found_in_iteration < 1
|
|
|
|
print(found)
|
|
|
|
--[[
|
|
For part 2, we use a recursive function to find how many bags a single bag can hold.
|
|
We use a table to provide memoization, gotta go fast.
|
|
--]]
|
|
local part2memo = {}
|
|
local function count_bags(color)
|
|
if part2memo[color] then return part2memo[color] end
|
|
local total = 0
|
|
for sub_color, count in pairs(contains[color]) do
|
|
-- Count the bag itself and all the bags it contains
|
|
total = total + count * (count_bags(sub_color) + 1)
|
|
end
|
|
part2memo[color] = total
|
|
return total
|
|
end
|
|
|
|
print(count_bags("shiny gold"))
|