advent-of-code-2019/day12.py

126 lines
2.8 KiB
Python

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