67 lines
1.5 KiB
Python
67 lines
1.5 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 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()
|