Optimize day 4 slightly

bram/rust/src/day_04/mod.rs

@@ -1,54 +1,119 @@
 // use crate::utils;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 pub fn answer(text : String) -> ( u32, u32 ) {
-    let mut paper_rolls : HashSet<( i32, i32 )> = text
-        .split_whitespace()
-        .enumerate()
-        .flat_map(
-            |(row, s)| s.chars().enumerate().filter_map(
-                move |(col, c)| if c == '@' { Some(( col as i32, row as i32 )) } else { None }
-            )
-        )
-        .collect();
+    let mut rolls = PaperRolls::from(&text);
+    let part_1 = rolls.remove();
 
-    let part_1 = paper_rolls
-        .iter()
-        .filter(|(x, y)| count_neighbours(&paper_rolls, *x, *y) < 5)
-        .count();
+    while rolls.remove() > 0 {};
 
-    let mut removes : HashSet<( i32, i32 )>;
-    let mut part_2 : u32 = 0;
-
-    loop {
-        (removes, paper_rolls) = paper_rolls
-            .iter()
-            .partition(|(x, y)| count_neighbours(&paper_rolls, *x, *y) < 5);
-
-        // println!("removes: ({} total, {} remaining) {:?}", removes.len(), paper_rolls.len(), removes);
-
-        if removes.len() == 0 {
-            break;
-        } else {
-            part_2 += removes.len() as u32;
-        }
-    }
-
-    ( part_1 as u32, part_2 )
+    ( part_1, rolls.rolls_moved )
 }
 
-fn count_neighbours(paper_rolls : &HashSet<( i32, i32 )>, x : i32, y : i32) -> u8 {
-    let mut total: u8 = 0;
+struct PaperRolls {
+    paper : HashMap<( u8, u8 ), u8>,
+    prev_move : Option<Vec<(u8, u8)>>,
+    rolls_moved : u32,
+}
+impl PaperRolls {
+    fn remove(&mut self) -> u32 {
+        let to_remove: Vec<(u8, u8)>;
+        
+        if let Some(prev_move) = &self.prev_move {
+            // We have already moved before. If we couldn't remove a roll last
+            // time, we won't be able to move it now unless it was a neighbour
+            // of a roll that was moved last turn.
+            let mut items = HashSet::new();
 
-    for nx in (x-1)..=(x+1) {
-        for ny in (y-1)..=(y+1) {
-            if paper_rolls.contains(&(nx, ny)) {
+            for (x, y) in prev_move {
+                let x1 = if x == &u8::MIN { u8::MIN } else { x - 1 };
+                let x2 = if x == &u8::MAX { u8::MAX } else { x + 1 };
+                let y1 = if y == &u8::MIN { u8::MIN } else { y - 1 };
+                let y2: u8 = if y == &u8::MAX { u8::MAX } else { y + 1 };
+
+                for nx in x1..=x2 {
+                    for ny in y1..=y2 {
+                        self.paper
+                            .entry((nx, ny))
+                            .and_modify(|e| {
+                                let x: u8 = *e;
+                                if x > 4 {
+                                    *e -= 1;
+                                } else {
+                                    items.insert((nx, ny));
+                                }
+                            });
+                    }
+                }
+            }
+
+            to_remove = items.into_iter().collect();
+        } else {
+            // This is (likely) the first time we remove rolls!
+            // Hence, do the less optimized option of finding all rolls
+            // with few neighbours directly in the HashMap.
+            to_remove = self.paper
+                .iter()
+                .filter(|(_, v)| **v < 4)
+                .map(|(k, _)| *k)
+                .collect();
+        }
+        
+        for r in &to_remove {
+            self.paper.remove(r);
+        }
+
+        let size: u32 = to_remove.len() as u32;
+
+        self.prev_move = Some(to_remove);
+        self.rolls_moved += size;
+        size
+    }
+    
+    fn from(s : &str) -> PaperRolls {
+        let rolls : HashSet<( u8, u8 )> = s
+            .split_whitespace()
+            .enumerate()
+            .flat_map(move |(row, s)|
+                s.chars().enumerate().filter_map(move |(col, c)|
+                    if c == '@' { Some(( col as u8, row as u8 )) } else { None }
+                )
+            )
+            .collect();
+
+
+        PaperRolls {
+            paper: rolls
+                .iter()
+                .map(|(x, y)|
+                    ((*x, *y), count_neighbours(&rolls, x, y))
+                )
+                .collect(),
+            prev_move : None,
+            rolls_moved: 0,
+        }
+    }
+}
+
+fn count_neighbours(rolls : &HashSet<(u8, u8)>, x : &u8, y : &u8) -> u8 {
+    let x1 : u8 = if x == &u8::MIN { u8::MIN } else { x - 1 };
+    let x2 : u8 = if x == &u8::MAX { u8::MAX } else { x + 1 };
+    let y1 : u8 = if y == &u8::MIN { u8::MIN } else { y - 1 };
+    let y2 : u8 = if y == &u8::MAX { u8::MAX } else { y + 1 };
+
+    let mut total = 0;
+
+    for nx in x1..=x2 {
+        for ny in y1..=y2 {
+            if x == &nx && y == &ny {
+                continue;
+            }
+
+            if rolls.contains(&(nx, ny)) {
                 total += 1;
             }
         }
     }
 
-    // println!("({}, {}) has {} neighbours - including itself.", x, y, total);
-
     total
 }

bram/rust/src/utils/mod.rs

@@ -40,9 +40,9 @@ pub fn str_to_i16(s : &str) -> Option<i16> {
 //     s.trim().to_string().parse::<u32>().ok()
 // }
     
-pub fn str_to_u64(s : &str) -> Option<u64> {
-    s.trim().to_string().parse::<u64>().ok()
-}
+// pub fn str_to_u64(s : &str) -> Option<u64> {
+//     s.trim().to_string().parse::<u64>().ok()
+// }
     
 pub fn str_to_u128(s : &str) -> Option<u128> {
     s.trim().to_string().parse::<u128>().ok()