Rewrote part 2 using Chinese Remainder theorem

aoc_2024/src/day14.rs

@@ -1,5 +1,3 @@
-use std::sync::OnceLock;
-
 use aoc_runner_derive::{aoc, aoc_generator};
 
 use nom::{
@@ -10,6 +8,8 @@ use nom::{
     sequence::{preceded, separated_pair},
 };
 
+use num_modular::ModularPow;
+
 type Input = Vec<((i64, i64), (i64, i64))>;
 
 #[aoc_generator(day14)]
@@ -54,43 +54,42 @@ fn part1(input: &Input) -> i64 {
 }
 
 #[aoc(day14, part2)]
-fn part2(input: &Input) -> usize {
+fn part2(input: &Input) -> i64 {
     let (w, h) = (101, 103);
 
-    let mut input = input.clone();
-
-    let mut iter = 0;
-
-    loop {
-        iter += 1;
-
-        let mut total = (0f64, 0f64);
-        input = input
-            .into_iter()
-            .map(|((px, py), (vx, vy))| {
-                let (px, py) = ((px + vx + w) % w, (py + vy + h) % h);
-                total.0 += px as f64;
-                total.1 += py as f64;
-                ((px, py), (vx, vy))
-            })
-            .collect();
-        let mean = (total.0 / input.len() as f64, total.1 / input.len() as f64);
-        let variance = input
-            .iter()
-            .map(|((px, py), _)| {
-                let dx = *px as f64 - mean.0;
-                let dy = *py as f64 - mean.1;
-                dx * dx + dy * dy
-            })
+    let mut offset_x = 0;
+    let mut min_var = f64::INFINITY;
+    for t in 0..w {
+        let values = input.iter().map(|((p, _), (v, _))| (p + t * (v + w)) % w);
+        let mean = values.clone().sum::<i64>() as f64 / input.len() as f64;
+        let variance = values
+            .map(|v| (v as f64 - mean) * (v as f64 - mean))
             .sum::<f64>()
             / input.len() as f64;
-
-        if variance < 1000.0 {
-            break;
+        if variance < min_var {
+            min_var = variance;
+            offset_x = t;
         }
     }
 
-    iter
+    let mut offset_y = 0;
+    let mut min_var = f64::INFINITY;
+    for t in 0..h {
+        let values = input.iter().map(|((_, p), (_, v))| (p + t * (v + h)) % h);
+        let mean = values.clone().sum::<i64>() as f64 / input.len() as f64;
+        let variance = values
+            .map(|v| (v as f64 - mean) * (v as f64 - mean))
+            .sum::<f64>()
+            / input.len() as f64;
+        if variance < min_var {
+            min_var = variance;
+            offset_y = t;
+        }
+    }
+
+    let invmod_x = (w as u64).powm((h - 2) as u64, &(h as u64)) as i64;
+    let invmod_y = (h as u64).powm((w - 2) as u64, &(w as u64)) as i64;
+    (offset_x * h * invmod_x + offset_y * w * invmod_y) % (w * h)
 }
 
 #[cfg(test)]