diff --git a/2020/d07/ex1/ex1.py b/2020/d07/ex1/ex1.py new file mode 100755 index 0000000..d92977c --- /dev/null +++ b/2020/d07/ex1/ex1.py @@ -0,0 +1,77 @@ +#!/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()