optimise

Author: vivek-26

aoc-2023/day-24/src/bin/part_two.rs

@@ -1,8 +1,8 @@
 use itertools::Itertools;
-use z3::ast::{Ast, Int};
+use z3::ast::{Ast, Int, Real};
 
 #[aoc_runner::main(24)]
-fn main(input: &str) -> i64 {
+fn main(input: &str) -> usize {
     let hailstones = input
         .split('\n')
         .map(|line| {
@@ -24,45 +24,32 @@ fn main(input: &str) -> i64 {
     solve(&hailstones)
 }
 
-fn solve(hailstones: &[((f64, f64, f64), (f64, f64, f64))]) -> i64 {
-    let cfg = z3::Config::new();
-    let ctx = z3::Context::new(&cfg);
+fn solve(hailstones: &[((f64, f64, f64), (f64, f64, f64))]) -> usize {
+    let ctx = z3::Context::new(&z3::Config::new());
     let s = z3::Solver::new(&ctx);
+    let [fx, fy, fz, fdx, fdy, fdz] =
+        ["fx", "fy", "fz", "fdx", "fdy", "fdz"].map(|v| Real::new_const(&ctx, v));
 
-    let (fx, fy, fz) = (
-        Int::new_const(&ctx, "fx"),
-        Int::new_const(&ctx, "fy"),
-        Int::new_const(&ctx, "fz"),
-    );
-    let (fdx, fdy, fdz) = (
-        Int::new_const(&ctx, "fdx"),
-        Int::new_const(&ctx, "fdy"),
-        Int::new_const(&ctx, "fdz"),
-    );
-
-    let zero = Int::from_i64(&ctx, 0);
+    let zero = Int::from_i64(&ctx, 0).to_real();
     for (i, &((x, y, z), (dx, dy, dz))) in hailstones.iter().enumerate() {
-        let (x, y, z) = (
-            Int::from_i64(&ctx, x as _),
-            Int::from_i64(&ctx, y as _),
-            Int::from_i64(&ctx, z as _),
-        );
-        let (dx, dy, dz) = (
-            Int::from_i64(&ctx, dx as _),
-            Int::from_i64(&ctx, dy as _),
-            Int::from_i64(&ctx, dz as _),
-        );
+        let [x, y, z, dx, dy, dz] =
+            [x, y, z, dx, dy, dz].map(|v| Int::from_i64(&ctx, v as _).to_real());
+
+        let t = Real::new_const(&ctx, format!("t{i}"));
 
-        let t = Int::new_const(&ctx, format!("t{i}"));
         s.assert(&t.ge(&zero));
         s.assert(&((&x + &dx * &t)._eq(&(&fx + &fdx * &t))));
         s.assert(&((&y + &dy * &t)._eq(&(&fy + &fdy * &t))));
         s.assert(&((&z + &dz * &t)._eq(&(&fz + &fdz * &t))));
     }
 
     assert_eq!(s.check(), z3::SatResult::Sat);
-    let model = s.get_model().unwrap();
-    let res = model.eval(&(&fx + &fy + &fz), true).unwrap();
 
-    res.as_i64().unwrap()
+    let res = s
+        .get_model()
+        .unwrap()
+        .eval(&(&fx + &fy + &fz), true)
+        .unwrap();
+
+    res.to_string().strip_suffix(".0").unwrap().parse().unwrap()
 }