optimize day 11 with linear time sum

README.md

@@ -27,19 +27,19 @@ Solutions for [Advent of Code](https://adventofcode.com/) in [Rust](https://www.
 
 | Day | Part 1 | Part 2 |
 | :---: | :---: | :---:  |
-| [Day 1](./src/bin/01.rs) | `29.6µs` | `38.5µs` |
-| [Day 2](./src/bin/02.rs) | `42.7µs` | `41.7µs` |
-| [Day 3](./src/bin/03.rs) | `88.4µs` | `100.0µs` |
-| [Day 4](./src/bin/04.rs) | `47.8µs` | `51.7µs` |
-| [Day 5](./src/bin/05.rs) | `22.1µs` | `24.5µs` |
-| [Day 6](./src/bin/06.rs) | `203.0ns` | `99.0ns` |
-| [Day 7](./src/bin/07.rs) | `101.7µs` | `95.7µs` |
-| [Day 8](./src/bin/08.rs) | `77.6µs` | `160.3µs` |
-| [Day 9](./src/bin/09.rs) | `47.7µs` | `48.6µs` |
-| [Day 10](./src/bin/10.rs) | `397.8µs` | `767.7µs` |
-| [Day 11](./src/bin/11.rs) | `28.7µs` | `28.6µs` |
-
-**Total: 2.24ms**
+| [Day 1](./src/bin/01.rs) | `32.6µs` | `35.8µs` |
+| [Day 2](./src/bin/02.rs) | `42.6µs` | `41.4µs` |
+| [Day 3](./src/bin/03.rs) | `84.2µs` | `99.7µs` |
+| [Day 4](./src/bin/04.rs) | `49.5µs` | `49.7µs` |
+| [Day 5](./src/bin/05.rs) | `20.7µs` | `24.1µs` |
+| [Day 6](./src/bin/06.rs) | `194.0ns` | `102.0ns` |
+| [Day 7](./src/bin/07.rs) | `96.3µs` | `92.1µs` |
+| [Day 8](./src/bin/08.rs) | `73.4µs` | `161.1µs` |
+| [Day 9](./src/bin/09.rs) | `51.8µs` | `47.9µs` |
+| [Day 10](./src/bin/10.rs) | `394.6µs` | `788.1µs` |
+| [Day 11](./src/bin/11.rs) | `16.4µs` | `15.8µs` |
+
+**Total: 2.22ms**
 <!--- benchmarking table --->
 
 ---

src/bin/11.rs

@@ -9,51 +9,67 @@ pub fn part_two(input: &str) -> Option<u64> {
 }
 
 pub fn compute_expansion(input: &str, expansion_factor: u32) -> Option<u64> {
-    let expansion_addition = expansion_factor.saturating_sub(1) as i32;
+    let expansion_addition = expansion_factor.saturating_sub(1) as usize;
+
+    // First compute the number of galaxies in each row and column.
+    let mut row_counts = Vec::with_capacity(input.len());
     let mut col_counts = Vec::new();
-    let mut galaxies = Vec::with_capacity(input.len());
-    let mut row_offsets = Vec::with_capacity(input.len());
-    let mut total_row_offset = 0;
-    for (row, line) in input.lines().enumerate() {
+    for line in input.lines() {
         if col_counts.is_empty() {
             col_counts.resize(line.len(), 0);
         }
         let mut row_count = 0;
         for (col, ch) in line.chars().enumerate() {
             if ch == '#' {
-                galaxies.push((row as i32, col as i32));
                 col_counts[col] += 1;
                 row_count += 1;
             }
         }
-        row_offsets.push(total_row_offset);
-        if row_count == 0 {
-            total_row_offset += expansion_addition;
-        }
-    }
-    let mut col_offsets = col_counts;
-    let mut total_col_offset = 0;
-    for val in &mut col_offsets {
-        let offset = total_col_offset;
-        if *val == 0 {
-            total_col_offset += expansion_addition;
-        }
-        *val = offset;
-    }
-    for galaxy in &mut galaxies {
-        galaxy.0 += row_offsets[galaxy.0 as usize];
-        galaxy.1 += col_offsets[galaxy.1 as usize];
+        row_counts.push(row_count);
     }
+
+    // Now compute the sum of pairwise row differences and column differences
+    // after expansion.
     Some(
-        galaxies
-            .iter()
-            .enumerate()
-            .flat_map(|(i, &g1)| galaxies.iter().skip(i + 1).map(move |&g2| (g1, g2)))
-            .map(|(g1, g2)| (g1.0.abs_diff(g2.0) + g1.1.abs_diff(g2.1)) as u64)
-            .sum(),
+        compute_diff_sum(&row_counts, expansion_addition)
+            + compute_diff_sum(&col_counts, expansion_addition),
     )
 }
 
+fn compute_diff_sum(counts: &[i32], expansion_addition: usize) -> u64 {
+    // This helper leverages an algebraic trick to turn the sum of distances
+    // into a linear operation:
+    // (x[i] - x[0]) + (x[i] - x[1]) + ... + (x[i] - x[i-1]) =
+    // i * x[i] - (x[0] + x[1] + ... + x[i - 1])
+    //
+    // Since we are iterating over galaxies computing the row difference
+    // and column difference according to that formula for every i from 1
+    // to N, we reduce the operation to computing a running sum and then
+    // subtracting i * x[i] - running_sum for every i, reducing a quadratic
+    // operation into a linear one.
+    //
+    // HOWEVER! This requires the values to be sorted. To do this, we
+    // reproduce the indexes by the row and column counts, which we
+    // needed in order to produce the expanded row/col indexes in the first
+    // place. So long as we iterate over those arrays in order, we are good!
+    let mut difference_sum = 0;
+    let mut total_sum = 0;
+    let mut offset = 0;
+    let mut galaxy_index = 0;
+    for (i, c) in counts.iter().copied().enumerate() {
+        let expanded_rowcol_index = i + offset;
+        for _ in 0..c {
+            difference_sum += galaxy_index * expanded_rowcol_index - total_sum;
+            total_sum += expanded_rowcol_index;
+            galaxy_index += 1;
+        }
+        if c == 0 {
+            offset += expansion_addition;
+        }
+    }
+    difference_sum as u64
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -71,3 +87,14 @@ mod tests {
         assert_eq!(compute_expansion(input, 100), Some(8410));
     }
 }
+
+// 5-2=3, 8-0=8, 3+8 = 11
+// 5-1=4, 9-8=1, 4+1 = 5
+// 5-0=5, 8-4=4, 5+4 = 9
+// 11 + 5 + 9 = 25
+
+// 3*5 - 3 = 15 - 3= 12
+// 3*8 - 13 = 24 - 13 = 11
+// 12 + 11 = 23
+
+// (x1 - x0) + (x2 - x1) + (x3 - x2) + (x4 - x3)