def test_apply_gravity():
moons = [
Moon(Point3D(-1, 0, 2)),
Moon(Point3D(2, -10, -7)),
Moon(Point3D(4, -8, 8)),
Moon(Point3D(3, 5, -1)),
]
moons = apply_gravity(moons)
assert moons[0].vel == (3, -1, -1)
def test_simulation():
moons = [
Moon(Point3D(-1, 0, 2)),
Moon(Point3D(2, -10, -7)),
Moon(Point3D(4, -8, 8)),
Moon(Point3D(3, 5, -1)),
]
for _ in range(10):
apply_gravity(moons)
apply_velocity(moons)
assert moons[0].pos == (2, 1, -3)
assert moons[0].vel == (-3, -2, 1)
assert calculate_energy(moons) == 179
def apply_gravity(moons: List[Moon]):
for moon in moons:
others = set(moons) - set([moon])
diff = [0, 0, 0]
for axis in range(3):
for other in others:
if moon.pos[axis] < other.pos[axis]:
diff[axis] += 1
if moon.pos[axis] > other.pos[axis]:
diff[axis] -= 1
moon.vel = Point3D(moon.vel.x + diff[0], moon.vel.y + diff[1],
moon.vel.z + diff[2])
return moons
def __init__(self, pos: Point3D, vel: Point3D = Point3D(0, 0, 0)):
self.pos = pos
self.vel = vel
def create_moons(positions: List[Point3D]):
return [Moon(Point3D(pos[0], pos[1], pos[2])) for pos in positions]
def apply_velocity(moons: List[Moon]):
for moon in moons:
moon.pos = Point3D(moon.pos.x + moon.vel.x, moon.pos.y + moon.vel.y,
moon.pos.z + moon.vel.z)
return moons
def test_apply_velocity():
moon = Moon(Point3D(-1, 0, 2), Point3D(3, -1, -1))
apply_velocity([moon])
assert moon.pos == (2, -1, 1)