mirror of
https://github.com/bertptrs/adventofcode.git
synced 2025-12-25 21:00:31 +01:00
70 lines
1.6 KiB
Python
70 lines
1.6 KiB
Python
import itertools
|
|
import math
|
|
import re
|
|
from typing import TextIO
|
|
|
|
import numpy # type: ignore
|
|
|
|
|
|
def read_moons(data: TextIO) -> numpy.array:
|
|
moons = []
|
|
|
|
for line in data:
|
|
moon = [int(x) for x in re.split(r"[^-0-9]+", line) if x]
|
|
moons.append(moon)
|
|
|
|
return numpy.array(moons)
|
|
|
|
|
|
def advance(pos: numpy.array, vel: numpy.array) -> None:
|
|
""" update pos and vel in place """
|
|
pos_prime = numpy.repeat(numpy.reshape(pos, (1, len(pos))), len(pos), axis=0).transpose() - pos
|
|
pos_prime = -numpy.sign(pos_prime)
|
|
vel += numpy.sum(pos_prime, axis=1)
|
|
pos += vel
|
|
|
|
|
|
def simulate_moons(moons: numpy.array, iterations: int) -> int:
|
|
moons = numpy.transpose(moons) # Transpose so we have rows of x, y, z
|
|
velocity = numpy.zeros_like(moons)
|
|
|
|
for _ in range(iterations):
|
|
for pos, vel in zip(moons, velocity):
|
|
advance(pos, vel)
|
|
|
|
potential = numpy.sum(numpy.abs(moons), axis=0)
|
|
kinetic = numpy.sum(numpy.abs(velocity), axis=0)
|
|
|
|
return int(numpy.sum(kinetic * potential))
|
|
|
|
|
|
def find_repetition(moons: numpy.array) -> int:
|
|
moons = numpy.transpose(moons)
|
|
velocity = numpy.zeros_like(moons)
|
|
|
|
needed = 1
|
|
|
|
for pos, vel in zip(moons, velocity):
|
|
pos_prime = numpy.copy(pos)
|
|
|
|
for i in itertools.count(1):
|
|
advance(pos, vel)
|
|
|
|
if (pos == pos_prime).all() and (vel == 0).all():
|
|
needed *= i // math.gcd(needed, i)
|
|
break
|
|
|
|
return needed
|
|
|
|
|
|
def part1(data: TextIO) -> int:
|
|
moons = read_moons(data)
|
|
|
|
return simulate_moons(moons, 1000)
|
|
|
|
|
|
def part2(data: TextIO) -> int:
|
|
moons = read_moons(data)
|
|
|
|
return find_repetition(moons)
|