Faster 2023 day 24 (part 1).

src/holt59/aoc/2023/day24.py

@@ -5,37 +5,32 @@ 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
-]
+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]
+for i1, (p1, v1) in enumerate(zip(positions, velocities)):
+    p, r = p1[:2], v1[:2]
 
-        rs = np.cross(r, s)
-        qpr = np.cross((q - p), r)
+    q, s = positions[i1 + 1 :, :2], velocities[i1 + 1 :, :2]
 
-        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
+    rs = np.cross(r, s)
 
-            if t >= 0 and u >= 0:
-                c = p + t * r
+    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()
 
-                if low <= c[0] <= high and low <= c[1] <= high:
-                    count += 1
 
 answer_1 = count
 print(f"answer 1 is {answer_1}")
@@ -58,8 +53,8 @@ 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]):
+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]