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

78 lines
1.8 KiB
Python
Raw Normal View History

2020-12-07 07:09:03 +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 reverse(graph: Graph) -> Graph:
reverse: Graph = defaultdict(set)
for color, contained in graph.items():
for col in contained:
reverse[col.color] |= {ColorInfo(-1, color)}
return reverse
def containing(graph: Graph, col: str) -> Set[ColorInfo]:
res = deepcopy(graph[col])
for contains in graph[col]:
res |= containing(graph, contains.color)
return res
def solve(raw: List[str]) -> int:
graph = to_graph(raw)
reverse_graph = reverse(graph)
return len(containing(reverse_graph, "shiny gold bag"))
def main() -> None:
input = [line.strip() for line in sys.stdin.readlines()]
print(solve(input))
if __name__ == "__main__":
main()