adventofcode2024/day13/sol.py

def parse(file):
    while True:
        a = file.readline()
        if not a:
            break
        b = file.readline()
        pr = file.readline()
        blank = file.readline()

        ax, ay = parse_numbers(a)
        bx, by = parse_numbers(b)
        px, py = parse_numbers(pr)
        yield ax,ay, bx,by, px,py

import re
def parse_numbers(s):
    return [int(x) for x in re.findall(r"\d+", s)]

from fractions import Fraction as F

import numpy
def solve(file):
    total = 0
    total2 = 0
    records = list(parse(file))
    for record in records:
        #print(record)
        ax, ay, bx, by, px, py = record

        # Part 1

        # Solve the linear equation:
        # k_a ax + k_b * bx = px
        # k_a ay + k_b * by = py

        M = numpy.array([[ax, bx], [ay, by]])
        V = numpy.array([px, py])

        sol = numpy.linalg.solve(M,V)
        k1 = int(round(sol[0]))
        k2 = int(round(sol[1]))

        if k1*ax + k2*bx == px and k1*ay + k2*by == py:
            #print("yes", k1, k2)
            total += 3*k1 + k2

        # Part 2
        qx = px + 10000000000000
        qy = py + 10000000000000

        V2 = numpy.array([qx, qy])
        sol = numpy.linalg.solve(M,V2)
        k1 = int(round(sol[0]))
        k2 = int(round(sol[1]))

        if k1*ax + k2*bx == qx and k1*ay + k2*by == qy:
            #print("yes2", k1, k2)
            total2 += 3*k1 + k2

    print(total)
    print(total2)


solve(open("sample1.in"))
solve(open("input"))