#!/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))