advent-of-code/src/holt59/aoc/2023/day24.py

import sys
from collections import defaultdict
from dataclasses import dataclass
import numpy as np

from sympy import solve, symbols


lines = sys.stdin.read().splitlines()

positions = [
    np.array([int(c) for c in line.split("@")[0].strip().split(", ")]) for line in lines
]
velocities = [
    np.array([int(c) for c in line.split("@")[1].strip().split(", ")]) for line in lines
]

# part 1
low, high = [7, 27] if len(positions) <= 10 else [200000000000000, 400000000000000]
count = 0
for i, (p1, v1) in enumerate(zip(positions, velocities)):
    for p2, v2 in zip(positions[i + 1:], velocities[i + 1:]):
        p, r = p1[:2], v1[:2]
        q, s = p2[:2], v2[:2]

        rs = np.cross(r, s)
        qpr = np.cross((q - p), r)

        if rs == 0:
            # assume there are no colinear lines
            continue
        else:
            assert rs != 0
            t = np.cross((q - p), s) / rs 
            u = np.cross((q - p), r) / rs

            if t >= 0 and u >= 0:
                c = p + t * r

                if low <= c[0] <= high and low <= c[1] <= high:
                    count += 1

answer_1 = count
print(f"answer 1 is {answer_1}")

# part 2
# equation
#   p1 + t1 * v1 == p0 + t1 * v0
#   p2 + t2 * v2 == p0 + t2 * v0
#   p3 + t3 * v3 == p0 + t3 * v0
#      ...
#   pn + tn * vn == p0 + tn * v0
# 

# we can solve with only 3 lines since each lines contains 3 
# equations (x / y / z), so 3 lines give 9 equations and 9
# variables: position (3), velocities (3) and times (3).
n = 3

x, y, z, vx, vy, vz, *ts = symbols('x y z vx vy vz ' + ' '.join(f't{i}' for i in range(n + 1)))
equations = []
for i, ti in zip(range(n), ts):
    for p, d, pi, di in zip((x, y, z), (vx, vy, vz), positions[i], velocities[i]):
        equations.append(p + ti * d - pi - ti * di)

r = solve(equations, [x, y, z, vx, vy, vz] + list(ts), dict=True)[0]

answer_2 = r[x] + r[y] + r[z]
print(f"answer 2 is {answer_2}")