### re: Advent of Code 2019 Solution Megathread - Day 12: The N-Body Problem VIEW POST

Part 1 in "modern" Python complete with tests. I need to think a bit about part 2. Can it be done with some form of integration?

``````from dataclasses import dataclass, field
from typing import List, Generator, TypeVar
import pytest

@dataclass(frozen=True)
class Vector:
x: int = 0
y: int = 0
z: int = 0

def magnitude(self):
return abs(self.x) + abs(self.y) + abs(self.z)

return Vector(self.x+other.x, self.y+other.y, self.z+other.z)

@dataclass(frozen=True)
class Moon:
pos: Vector
velocity: Vector = field(default_factory=Vector)

def energy(self):
return self.pos.magnitude() * self.velocity.magnitude()

"""
Load multi-line `text` into a collection of Moons
"""
def vector(line):
parts = line.strip("<>").split(",")
return Vector(*[int(p.strip()[2:]) for p in parts])

return [Moon(pos=vector(line)) for line in text.strip().splitlines()]

def simulate(moons: List[Moon]) -> List[Moon]:
"""
Simulate the motion of `moons`, returning a new collection of Moons
"""
def compare_axis(x, y):
if x < y: return 1
elif x == y: return 0
else: return -1

def gravity_vector(a: Vector, b: Vector) -> Vector:
return Vector(
x = compare_axis(a.x, b.x),
y = compare_axis(a.y, b.y),
z = compare_axis(a.z, b.z)
)

def simulate_moon(m: Moon) -> Moon:
gravity = Vector()
for n in moons:
gravity += gravity_vector(m.pos, n.pos)
velocity = m.velocity + gravity
return Moon(pos=m.pos + velocity, velocity=velocity)

return [simulate_moon(m) for m in moons]

def simulate_seq(moons: List[Moon]) -> Generator[List[Moon], None, None]:
"""
Generate a sequence of simulated states for `moons`
"""
m = moons
while True:
m = simulate(m)
yield m

T = TypeVar("T")
def after(n: int, seq: Generator[T, None, None]) -> T:
"""
Get the `n`th value from `seq`
"""
for x in range(n):
r = next(seq)
return r

def total_energy(moons: List[Moon]) -> int:
return sum(m.energy() for m in moons)

"""
Verify parsing text into Moons
"""
<x=2, y=-10, z=-7>
<x=4, y=-8, z=8>
<x=3, y=5, z=-1>""")
assert len(r) == 4
assert r[0] == Moon(pos=Vector(-1, 0, 2))
assert r[1] == Moon(pos=Vector(2,-10,-7))
assert r[2] == Moon(pos=Vector(4,-8,8))
assert r[3] == Moon(pos=Vector(3,5,-1))

EXAMPLE_MOONS_1 = [
Moon(pos=Vector(-1,  0, 2)),
Moon(pos=Vector( 2,-10,-7)),
Moon(pos=Vector( 4, -8, 8)),
Moon(pos=Vector( 3,  5,-1))
]

EXAMPLE_MOONS_2 = [
Moon(pos=Vector(-8,-10,  0)),
Moon(pos=Vector( 5,  5, 10)),
Moon(pos=Vector( 2, -7,  3)),
Moon(pos=Vector( 9, -8, -3))
]

@pytest.mark.parametrize("moons,steps,expect", [
(EXAMPLE_MOONS_1, 1, [
Moon(pos=Vector(2,-1, 1), velocity=Vector( 3,-1,-1)),
Moon(pos=Vector(3,-7,-4), velocity=Vector( 1, 3, 3)),
Moon(pos=Vector(1,-7, 5), velocity=Vector(-3, 1,-3)),
Moon(pos=Vector(2, 2, 0), velocity=Vector(-1,-3, 1))
]),
(EXAMPLE_MOONS_1, 2, [
Moon(pos=Vector(5,-3,-1), velocity=Vector( 3,-2,-2)),
Moon(pos=Vector(1,-2, 2), velocity=Vector(-2, 5, 6)),
Moon(pos=Vector(1,-4,-1), velocity=Vector( 0, 3,-6)),
Moon(pos=Vector(1,-4, 2), velocity=Vector(-1,-6, 2))
]),
(EXAMPLE_MOONS_1, 10, [
Moon(pos=Vector(2, 1,-3), velocity=Vector(-3,-2, 1)),
Moon(pos=Vector(1,-8, 0), velocity=Vector(-1, 1, 3)),
Moon(pos=Vector(3,-6, 1), velocity=Vector( 3, 2,-3)),
Moon(pos=Vector(2, 0, 4), velocity=Vector( 1,-1,-1))
])
])
def test_simulate_one_step(moons, steps, expect):
"""
Verify moon simulation
"""
assert after(steps, simulate_seq(moons)) == expect

@pytest.mark.parametrize("moons,steps,energy", [
(EXAMPLE_MOONS_1, 10, 179),
(EXAMPLE_MOONS_2, 100, 1940)
])
def test_energy(moons, steps, energy):
final = after(steps, simulate_seq(moons))
assert total_energy(final) == energy

def part1(moons):
final = after(1000, simulate_seq(moons))
print(f"Part 1 : {total_energy(final)}")

if __name__ == "__main__":
with open("input.txt", "rt") as f: