advent/y2019/p12.py

import functools
import itertools
import re
import sys
from math import gcd


def calculate_energy(moons, vels):
    return sum(
        sum(abs(p) for p in moon) * sum(abs(v) for v in vel)
        for moon, vel in zip(moons, vels)
    )


def simulate(text, axes=slice(0, 3), lim=None):
    g = (int(n) for n in re.findall(r'-?\d+', text))
    moons = [point[axes] for point in zip(*[g] * 3)]
    vels = [[0, 0, 0] for m in moons]

    seen = set()
    for step in itertools.count():
        if lim is not None and step >= lim:
            break

        for (i, A), (j, B) in itertools.combinations(enumerate(moons), 2):
            dV = [max(a < b, 0) or -max(b < a, 0) for a, b in zip(A, B)]
            vels[i] = [v + d for v, d in zip(vels[i], dV)]
            vels[j] = [v - d for v, d in zip(vels[j], dV)]

        for i, moon in enumerate(moons):
            moons[i] = [p + v for p, v in zip(moons[i], vels[i])]

        key = str(moons) + str(vels)
        if key in seen:
            break
        seen.add(key)

    return step, calculate_energy(moons, vels)


text = sys.stdin.read()

_, energy = simulate(text, lim=1000)
print(energy)

nX, _ = simulate(text, axes=slice(0, 1))
nY, _ = simulate(text, axes=slice(1, 2))
nZ, _ = simulate(text, axes=slice(2, 3))
lcm = functools.reduce(lambda a, b: a * b // gcd(a, b), [nX, nY, nZ])
print(lcm)