import operator import sys from functools import reduce from typing import Literal, TypeAlias, cast CubeType: TypeAlias = Literal["red", "blue", "green"] MAX_CUBES: dict[CubeType, int] = {"red": 12, "green": 13, "blue": 14} # parse games lines = sys.stdin.read().splitlines() games: dict[int, list[dict[CubeType, int]]] = {} for line in filter(bool, lines): id_part, sets_part = line.split(":") games[int(id_part.split(" ")[-1])] = [ { cast(CubeType, s[1]): int(s[0]) for cube_draw in cube_set_s.strip().split(", ") if (s := cube_draw.split(" ")) } for cube_set_s in sets_part.strip().split(";") ] # part 1 answer_1 = sum( id for id, set_of_cubes in games.items() if all( n_cubes <= MAX_CUBES[cube] for cube_set in set_of_cubes for cube, n_cubes in cube_set.items() ) ) print(f"answer 1 is {answer_1}") # part 2 answer_2 = sum( reduce( operator.mul, (max(cube_set.get(cube, 0) for cube_set in set_of_cubes) for cube in MAX_CUBES), ) for set_of_cubes in games.values() ) print(f"answer 2 is {answer_2}")