2020: d07: ex2: add solution

This commit is contained in:
Bruno BELANYI 2020-12-07 07:09:28 +01:00
parent e28d9faeef
commit d3d35e2001

66
2020/d07/ex2/ex2.py Executable file
View file

@ -0,0 +1,66 @@
#!/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()