105 lines
2.4 KiB
Python
105 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from itertools import combinations, starmap
|
|
from math import gcd
|
|
from functools import reduce
|
|
|
|
from lib import vector
|
|
|
|
def lcm(a, b):
|
|
return (a * b) // gcd(a, b)
|
|
|
|
def cmp(x, y):
|
|
return (x < y) - (x > y)
|
|
|
|
class Moon:
|
|
def __init__(self, pos, vel=None):
|
|
self.pos = vector(pos)
|
|
if vel == None:
|
|
vel = vector([0]*len(self.pos))
|
|
self.vel = vector(vel)
|
|
|
|
def restrict(self, i):
|
|
return Moon((self.pos[i],), (self.vel[i],))
|
|
|
|
def apply_gravity(self, b):
|
|
self.vel += starmap(cmp, zip(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):
|
|
parse = lambda line : Moon(int(x[2:]) for x in line[1:-1].split(', '))
|
|
return System(map(parse, puzzle_input.split('\n')))
|
|
|
|
def partI(moons):
|
|
return simulate(moons, 1000).energy()
|
|
|
|
def partII(moons):
|
|
return int(find_period(moons))
|
|
|