advent-of-code-2019/day12.py

#!/usr/bin/env python3

from itertools import combinations
from math import gcd
from functools import reduce

def lcm(a, b):
    return (a * b) // gcd(a, b)

def comp(a, b):
    def f(arg):
        x, y = arg
        if x < y:
            return 1
        elif x == y:
            return 0
        return -1
    return Vector(map(f, zip(a, b)))

def parse(line):
    coords = (int(x[2:]) for x in line[1:-1].split(', '))
    return Moon(coords)

class Vector(tuple):

    def __add__(self, other):
        return Vector(map(sum, zip(self, other)))

    def __neg__(self):
        return Vector((-x for x in self))

    def __sub__(self, other):
        return self + (- other)

    def __repr__(self):
        return super(Vector, self).__repr__()

class Moon:
    def __init__(self, pos, vel=None):
        self.pos = Vector(pos)
        if vel == None:
            self.vel = Vector([0]*len(self.pos))
        else:
            self.vel = Vector(vel)

    def restrict(self, i):
        return Moon((self.pos[i],), (self.vel[i],))

    def apply_gravity(self, b):
        self.vel += comp(self.pos, b.pos)

    def apply_velocity(self):
        self.pos += self.vel

    def copy(self):
        return Moon(self.pos, self.vel)

    def __len__(self):
        return len(self.pos)

    def energy(self):
        return sum(map(abs, self.pos)) * sum(map(abs, self.vel))

    def __eq__(self, other):
        return self.pos == other.pos and self.vel == other.vel

    def __ne__(self, other):
        return not self == other

    def __repr__(self):
        return "<{}>, <{}>".format(', '.join(map(str, self.pos)),
                                   ', '.join(map(str, self.vel)))


class System(tuple):

    def dimension(self):
        return min(len(m) for m in self)

    def restrict(self, i):
        return System(m.restrict(i) for m in self)

    def energy(self):
        return sum(m.energy() for m in self)

    def __repr__(self):
        return '\n'.join(map(str, self))

    def __next__(self):
        for a, b in combinations(self, 2):
            a.apply_gravity(b)
            b.apply_gravity(a)
        for m in self:
            m.apply_velocity()

    def copy(self):
        return System(m.copy() for m in self)


def simulate(sys, steps):
    sys = sys.copy()
    for i in range(steps):
        next(sys)
    return sys

def find_period(sys):
    if sys.dimension() == 1:
        future = sys.copy()
        next(future)
        step = 1
        while future != sys:
            next(future)
            step += 1
        return step
    return reduce(lcm, (find_period(sys.restrict(i))
                        for i in range(sys.dimension())))
def preproc(puzzle_input):
    return System(map(parse, puzzle_input.split('\n')))

def partI(moons):
    return simulate(moons, 1000).energy()

def partII(moons):
    return int(find_period(moons))