day12.py

#!/usr/bin/env python3
from dataclasses import dataclass, field
from functools import reduce
from itertools import combinations
from math import gcd
from typing import List, Set
import copy


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

@dataclass
class Planet:
    position: int = 0
    velocity: int = 0

    def apply_velocity(self):
        self.position += self.velocity

    def apply_gravity(self, other):
        if self.position > other.position:
            self.velocity -= 1
            other.velocity += 1
        elif self.position < other.position:
            self.velocity += 1
            other.velocity -= 1

@dataclass
class Dimension:
    points: List[Planet] = field(default_factory=list)

    def step(self):
        for p1, p2 in combinations(self.points, 2):
            p1.apply_gravity(p2)

        for p in self.points:
            p.apply_velocity()

        return self

    def period(self):
        dimension = copy.deepcopy(self)

        n = 1
        while dimension.step() != self:
            n += 1

        return n

@dataclass
class System:
    x: Dimension = field(default_factory=Dimension)
    y: Dimension = field(default_factory=Dimension)
    z: Dimension = field(default_factory=Dimension)

    def step(self):
        self.x.step()
        self.y.step()
        self.z.step()

    def period(self):
        x_period = self.x.period()
        y_period = self.y.period()
        z_period = self.z.period()
        return lcm([x_period, y_period, z_period])

    @staticmethod
    def _energy(x,y,z):
        potential = abs(x.position) + abs(y.position) + abs(z.position)
        kinetic   = abs(x.velocity) + abs(y.velocity) + abs(z.velocity)
        return potential * kinetic

    def energy(self):
        return sum(self._energy(x,y,z) for (x,y,z) in zip(self.x.points,self.y.points,self.z.points))

def load_planets():
    system = System()
    with open("day12.txt", "r") as f:
        for li in f.readlines():
            li = li.strip("<>\r\n ")
            planet = Planet()
            items = li.split(",")
            for item in items:
                k, v = item.split("=")
                planet = Planet(position=int(v))
                getattr(system, k.strip()).points.append(planet)
    return system

print("Part 1:")

system = load_planets()

for _ in range(1000):
    system.step()

print(system.energy())

print("Part 2:")

print(system.period())