import operator import sys from dataclasses import dataclass from functools import reduce from typing import Literal, TypeAlias, cast Cube: TypeAlias = Literal["red", "blue", "green"] @dataclass(frozen=True) class Game: id: int set_of_cubes: list[dict[Cube, int]] lines = sys.stdin.read().split("\n") games: list[Game] = [] for line in filter(bool, lines): id_part, sets_part = line.split(":") games.append( Game( id=int(id_part.split(" ")[-1]), set_of_cubes=[ { cast(Cube, 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 MAX_CUBES: dict[Cube, int] = {"red": 12, "green": 13, "blue": 14} answer_1 = sum( game.id for game in games if all( n_cubes <= MAX_CUBES[cube] for cube_set in game.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 game.set_of_cubes) for cube in MAX_CUBES ), ) for game in games ) print(f"answer 2 is {answer_2}")