Implement 2019 day 12

2019/aoc2019/day12.py

@@ -0,0 +1,69 @@
+import itertools
+import math
+import re
+from typing import TextIO
+
+import numpy  # type: ignore
+
+
+def read_moons(data: TextIO) -> numpy.array:
+    moons = []
+
+    for line in data:
+        moon = [int(x) for x in re.split(r"[^-0-9]+", line) if x]
+        moons.append(moon)
+
+    return numpy.array(moons)
+
+
+def advance(pos: numpy.array, vel: numpy.array) -> None:
+    """ update pos and vel in place """
+    pos_prime = numpy.repeat(numpy.reshape(pos, (1, len(pos))), len(pos), axis=0).transpose() - pos
+    pos_prime = -numpy.sign(pos_prime)
+    vel += numpy.sum(pos_prime, axis=1)
+    pos += vel
+
+
+def simulate_moons(moons: numpy.array, iterations: int) -> int:
+    moons = numpy.transpose(moons)  # Transpose so we have rows of x, y, z
+    velocity = numpy.zeros_like(moons)
+
+    for _ in range(iterations):
+        for pos, vel in zip(moons, velocity):
+            advance(pos, vel)
+
+    potential = numpy.sum(numpy.abs(moons), axis=0)
+    kinetic = numpy.sum(numpy.abs(velocity), axis=0)
+
+    return int(numpy.sum(kinetic * potential))
+
+
+def find_repetition(moons: numpy.array) -> int:
+    moons = numpy.transpose(moons)
+    velocity = numpy.zeros_like(moons)
+
+    needed = 1
+
+    for pos, vel in zip(moons, velocity):
+        pos_prime = numpy.copy(pos)
+
+        for i in itertools.count(1):
+            advance(pos, vel)
+
+            if (pos == pos_prime).all() and (vel == 0).all():
+                needed *= i // math.gcd(needed, i)
+                break
+
+    return needed
+
+
+def part1(data: TextIO) -> int:
+    moons = read_moons(data)
+
+    return simulate_moons(moons, 1000)
+
+
+def part2(data: TextIO) -> int:
+    moons = read_moons(data)
+
+    return find_repetition(moons)