advent-of-code/2020/d07/ex2/ex2.py

67 lines
1.5 KiB
Python
Raw Normal View History

2020-12-07 07:09:28 +01:00
#!/usr/bin/env python
import re
import sys
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass
from typing import Dict, List, Set, Tuple
@dataclass(eq=True, frozen=True) # Hashable
class ColorInfo:
num: int
color: str
Graph = Dict[str, Set[ColorInfo]]
def extract_info(line: str) -> Tuple[str, Set[ColorInfo]]:
color_pattern = re.compile("(.*) contain ")
match = color_pattern.search(line)
assert match is not None
color = match.group(1).replace("bags", "bag")
line = line[match.end() : -1] # Remove period at end of line
if line == "no other bags":
return color, set()
colors: Set[ColorInfo] = set()
pattern = re.compile("([0-9]+) (.*)")
for col in line.split(", "):
match = pattern.search(col)
assert match is not None
colors |= {
ColorInfo(int(match.group(1)), match.group(2).replace("bags", "bag"))
}
return color, colors
def to_graph(raw: List[str]) -> Graph:
return {color: inside for color, inside in map(extract_info, raw)}
def num_bags(graph: Graph, col: str) -> int:
return sum(
contained.num * (1 + num_bags(graph, contained.color))
for contained in graph[col]
)
def solve(raw: List[str]) -> int:
graph = to_graph(raw)
return num_bags(graph, "shiny gold bag")
def main() -> None:
input = [line.strip() for line in sys.stdin.readlines()]
print(solve(input))
if __name__ == "__main__":
main()