aoc2024day13.py

# === Parsing ===
from fractions import Fraction as F
from itertools import batched
import re

BUTTON_RE = re.compile(r"Button .: X\+(\d+), Y\+(\d+)")
PRIZE_RE = re.compile(r"Prize: X=(\d+), Y=(\d+)")

games = []
with open("input.txt", "r") as f:
    for batch_lines in batched(f, 4):
        button_a, button_b, prize, *_ = batch_lines
        ax, ay = BUTTON_RE.match(button_a).groups()
        bx, by = BUTTON_RE.match(button_b).groups()
        px, py = PRIZE_RE.match(prize).groups()
        games.append(((F(ax), F(ay)), (F(bx), F(by)), (F(px), F(py))))

print(games)


# === Part 1 ===
token_cost = 0
for (ax, ay), (bx, by), (px, py) in games:
    B = (py - px * ay / ax) / (by - bx * ay / ax)
    A = (px - B * bx) / ax
    print(A, B)
    if A == int(A) and B == int(B):
        token_cost += 3 * A + B
print(token_cost)


# === Part 2 ===
P_OFFSET = 10_000_000_000_000  # 10 trillion
token_cost = 0
for (ax, ay), (bx, by), (px, py) in games:
    px, py = P_OFFSET + px, P_OFFSET + py
    B = (py - px * ay / ax) / (by - bx * ay / ax)
    A = (px - B * bx) / ax
    print(A, B)
    if A == int(A) and B == int(B):
        token_cost += 3 * A + B
print(token_cost)