78 lines
1.8 KiB
Python
78 lines
1.8 KiB
Python
|
#!/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()
|