2023: d24: ex2: add solution
This commit is contained in:
parent
0c23bd9eee
commit
b004a8fae7
65
2023/d24/ex2/ex2.py
Executable file
65
2023/d24/ex2/ex2.py
Executable file
|
@ -0,0 +1,65 @@
|
|||
#!/usr/bin/env python
|
||||
|
||||
import sys
|
||||
from typing import NamedTuple
|
||||
|
||||
import z3
|
||||
|
||||
|
||||
class Point(NamedTuple):
|
||||
x: int
|
||||
y: int
|
||||
z: int
|
||||
|
||||
|
||||
class HailStone(NamedTuple):
|
||||
pos: Point
|
||||
vel: Point
|
||||
|
||||
|
||||
def solve(input: list[str]) -> int:
|
||||
def parse_line(line: str) -> HailStone:
|
||||
pos, vel = line.split(" @ ")
|
||||
return HailStone(
|
||||
Point(*map(int, pos.split(", "))), Point(*map(int, vel.split(", ")))
|
||||
)
|
||||
|
||||
def parse(input: list[str]) -> list[HailStone]:
|
||||
return [parse_line(line) for line in input]
|
||||
|
||||
# Bringing out the big guns to solve this
|
||||
def run_z3(hailstones: list[HailStone]) -> HailStone:
|
||||
solver = z3.Solver()
|
||||
|
||||
px, py, pz = (z3.Int(f"stone_p{coord}") for coord in ("x", "y", "z"))
|
||||
vx, vy, vz = (z3.Int(f"stone_v{coord}") for coord in ("x", "y", "z"))
|
||||
|
||||
for i, other in enumerate(hailstones):
|
||||
t = z3.Int(f"t_{i}")
|
||||
# Must be in the future
|
||||
solver.add(t >= 0)
|
||||
# Must hit the hailstone
|
||||
solver.add((px + vx * t) == (other.pos.x + other.vel.x * t))
|
||||
solver.add((py + vy * t) == (other.pos.y + other.vel.y * t))
|
||||
solver.add((pz + vz * t) == (other.pos.z + other.vel.z * t))
|
||||
|
||||
assert solver.check() == z3.sat # Sanity check
|
||||
|
||||
model = solver.model()
|
||||
return HailStone(
|
||||
Point(*(model.eval(v).as_long() for v in (px, py, pz))), # type: ignore
|
||||
Point(*(model.eval(v).as_long() for v in (vx, vy, vz))), # type: ignore
|
||||
)
|
||||
|
||||
hailstones = parse(input)
|
||||
stone = run_z3(hailstones)
|
||||
return sum(stone.pos)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
input = sys.stdin.read().splitlines()
|
||||
print(solve(input))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
Loading…
Reference in a new issue