69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
from sympy import solve, symbols
|
|
|
|
from ..base import BaseSolver
|
|
|
|
|
|
class Solver(BaseSolver):
|
|
def solve(self, input: str) -> Iterator[Any]:
|
|
lines = input.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 i1, (p1, v1) in enumerate(zip(positions, velocities)):
|
|
p, r = p1[:2], v1[:2]
|
|
|
|
q, s = positions[i1 + 1 :, :2], velocities[i1 + 1 :, :2]
|
|
|
|
rs = np.cross(r, s)
|
|
|
|
q, s, rs = q[m := (rs != 0)], s[m], rs[m]
|
|
t = np.cross((q - p), s) / rs
|
|
u = np.cross((q - p), r) / rs
|
|
|
|
t, u = t[m := ((t >= 0) & (u >= 0))], u[m]
|
|
c = p + np.expand_dims(t, 1) * r
|
|
count += np.all((low <= c) & (c <= high), axis=1).sum()
|
|
|
|
yield count
|
|
|
|
# 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 i1, ti in zip(range(n), ts):
|
|
for p, d, pi, di in zip(
|
|
(x, y, z), (vx, vy, vz), positions[i1], velocities[i1]
|
|
):
|
|
equations.append(p + ti * d - pi - ti * di)
|
|
|
|
r = solve(equations, [x, y, z, vx, vy, vz] + list(ts), dict=True)[0]
|
|
yield r[x] + r[y] + r[z]
|