126 lines
3.7 KiB
Python
126 lines
3.7 KiB
Python
import sys
|
|
import json
|
|
import itertools
|
|
|
|
import pydantic
|
|
import sympy
|
|
|
|
from typing import Iterable, TextIO
|
|
|
|
|
|
class Triplet(pydantic.BaseModel):
|
|
x: int
|
|
y: int
|
|
z: int
|
|
|
|
|
|
class Hailstone(pydantic.BaseModel):
|
|
position: Triplet
|
|
velocity: Triplet
|
|
|
|
|
|
def parseInputFile(inputFile: TextIO) -> list[Hailstone]:
|
|
inputData = json.load(inputFile)
|
|
return [Hailstone(**datum) for datum in inputData]
|
|
|
|
|
|
def part1(hailstones: Iterable[Hailstone]) -> int:
|
|
total = 0
|
|
for hailstone1, hailstone2 in itertools.combinations(hailstones, 2):
|
|
t = sympy.var("t") # type: ignore
|
|
x, y = sympy.var("x y") # type: ignore
|
|
hailstone1XEq = sympy.Eq(
|
|
hailstone1.position.x + hailstone1.velocity.x * t, x # type: ignore
|
|
)
|
|
|
|
hailstone1YEq = sympy.Eq(
|
|
hailstone1.position.y + hailstone1.velocity.y * t, y # type: ignore
|
|
)
|
|
|
|
hailstone2XEq = sympy.Eq(
|
|
hailstone2.position.x + hailstone2.velocity.x * t, x # type: ignore
|
|
)
|
|
|
|
hailstone2YEq = sympy.Eq(
|
|
hailstone2.position.y + hailstone2.velocity.y * t, y # type: ignore
|
|
)
|
|
|
|
# These expressions only have one solution, and we can guarantee that by their algebraic form
|
|
hailstone1Eq = hailstone1XEq.subs({t: sympy.solve(hailstone1YEq, t)[0]})
|
|
hailstone2Eq = hailstone2XEq.subs({t: sympy.solve(hailstone2YEq, t)[0]})
|
|
|
|
solns = sympy.solve(
|
|
[
|
|
hailstone1Eq,
|
|
hailstone2Eq,
|
|
],
|
|
x,
|
|
y,
|
|
dict=True,
|
|
)
|
|
if not solns:
|
|
continue
|
|
elif len(solns) != 1:
|
|
# Should never happen if they're linear
|
|
raise ValueError(
|
|
f"More than one solution for hailstones {hailstone1} {hailstone2}"
|
|
)
|
|
|
|
xSoln = solns[0][x]
|
|
ySoln = solns[0][y]
|
|
halestone1Time = sympy.solve(hailstone1XEq)[0][t].subs({x: xSoln})
|
|
halestone2Time = sympy.solve(hailstone2XEq)[0][t].subs({x: xSoln})
|
|
if (
|
|
halestone1Time < 0
|
|
or halestone2Time < 0
|
|
or not (200000000000000 <= xSoln <= 400000000000000)
|
|
or not (200000000000000 <= ySoln <= 400000000000000)
|
|
):
|
|
continue
|
|
|
|
total += 1
|
|
|
|
return total
|
|
|
|
|
|
def part2(hailstones: Iterable[Hailstone]) -> int:
|
|
equations = []
|
|
timeVariables = []
|
|
x0, y0, z0 = sympy.var("x y z") # type: ignore
|
|
vx, vy, vz = sympy.var("vx vy vz") # type: ignore
|
|
|
|
# We have 6 unknown variables for our velocity, and one variable (t_n) for every three equations.
|
|
# With 4 hailstones, we will have 12 equations, and 10 unknowns, which is more than sufficient
|
|
for i, hailstone in enumerate(itertools.islice(hailstones, 4), 1):
|
|
timeVar = sympy.var(f"t{i}")
|
|
localEquations = [
|
|
x0 + timeVar * vx - (hailstone.position.x + timeVar * hailstone.velocity.x), # type: ignore
|
|
y0 + timeVar * vy - (hailstone.position.y + timeVar * hailstone.velocity.y), # type: ignore
|
|
z0 + timeVar * vz - (hailstone.position.z + timeVar * hailstone.velocity.z), # type: ignore
|
|
]
|
|
|
|
timeVariables.append(timeVar)
|
|
equations.extend(localEquations)
|
|
|
|
res = sympy.solve(equations, *[x0, y0, z0, vx, vy, vz, *timeVariables], dict=True)
|
|
if len(res) != 1:
|
|
raise ValueError("more than one unique solution")
|
|
|
|
return res[0][x0] + res[0][y0] + res[0][z0]
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) != 2:
|
|
print(f"Usage: {sys.argv[0]} inputFile", sys.stderr)
|
|
sys.exit(1)
|
|
|
|
with open(sys.argv[1]) as inputFile:
|
|
inputData = parseInputFile(inputFile)
|
|
|
|
print("Part 2: ", part2(inputData))
|
|
print("Part 1: ", part1(inputData))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|