advent-of-code/2023/python/day02.py

import shared
from collections import defaultdict


def split_games(row):
    game, draws = row.split(":")
    game = int(game.split(" ")[-1])
    colors = defaultdict(list)
    for draw in draws.split(";"):
        for cube in draw.split(","):
            cube = cube.split()
            count, color = cube
            colors[color].append(int(count))
    return game, colors


MAXES = {
    "red": 12,
    "green": 13,
    "blue": 14,
}


# @shared.profile
def part1(rows):
    applicable = []
    for row in rows:
        ok = True
        game, colors = split_games(row)
        for color, _max in MAXES.items():
            bag_max = max(colors[color])
            if bag_max > _max:
                ok = False
                break
        if ok:
            applicable.append(game)
    print(sum(applicable))


# @shared.profile
def part2(rows):
    total = 0
    for row in rows:
        _, colors = split_games(row)
        power = 1
        for vals in colors.values():
            power *= max(vals)
        total += power
    print(total)


def main():
    rows = [row for row in shared.load_rows(2)]
    with shared.elapsed_timer() as elapsed:
        part1(rows)
        print("🕒", elapsed())

    rows = [row for row in shared.load_rows(2, True)]
    with shared.elapsed_timer() as elapsed:
        part2(rows)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()