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)

2019/tests/test_day12.py

@@ -0,0 +1,25 @@
+import pytest
+import numpy  # type: ignore
+
+from aoc2019.day12 import simulate_moons, find_repetition
+
+SAMPLES = [
+    numpy.array([[-1, 0, 2], [2, -10, -7], [4, -8, 8], [3, 5, -1]]),
+    numpy.array([[-8, -10, 0], [5, 5, 10], [2, -7, 3], [9, -8, -3]]),
+]
+
+
+@pytest.mark.parametrize('iterations,result,moons', [
+    (10, 179, numpy.copy(SAMPLES[0])),
+    (100, 1940, numpy.copy(SAMPLES[1])),
+])
+def test_kinetic_energy(moons: numpy.array, iterations: int, result: int) -> None:
+    assert simulate_moons(moons, iterations) == result
+
+
+@pytest.mark.parametrize('outcome,moons', [
+    (2772, numpy.copy(SAMPLES[0])),
+    (4686774924, numpy.copy(SAMPLES[1])),
+])
+def test_repetition(moons: numpy.array, outcome: int) -> None:
+    assert find_repetition(moons) == outcome