advent-of-code/2023/day2.py

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}")