Optimize day 14 with bitwise operations

Author: SkiFire13

src/day14.rs

@@ -1,120 +1,136 @@
 #[allow(unused_imports)]
 use super::prelude::*;
-type Input = Grid<u8>;
+type Input = (Box<[u128]>, Box<[u128]>);
 
 pub fn input_generator(input: &str) -> Input {
-    Grid::from_input_chars(input, |c, _, _| c as u8)
-}
+    let (mut rocks, mut walls) = (vec![0], vec![!0]);
+    let (mut rock_layer, mut wall_layer) = (0, 1);
 
-pub fn part1(input: &Input) -> usize {
-    let mut input = input.clone();
-    for y in 1..input.h() {
-        for x in 0..input.w() {
-            let mut y = y;
-            while y > 0 && input[(x, y)] == b'O' && input[(x, y - 1)] == b'.' {
-                input[(x, y)] = b'.';
-                input[(x, y - 1)] = b'O';
-                y -= 1;
+    for b in input.bytes() {
+        (rock_layer, wall_layer) = (rock_layer << 1, wall_layer << 1);
+        match b {
+            b'\n' => {
+                rocks.push(rock_layer);
+                walls.push(wall_layer | 1);
+                (rock_layer, wall_layer) = (0, 1);
             }
+            b'O' => rock_layer |= 1,
+            b'#' => wall_layer |= 1,
+            b'.' => {}
+            _ => panic!(),
         }
     }
 
-    let mut tot = 0;
+    rocks.push(0);
+    walls.push(!0);
 
-    for y in 0..input.h() {
-        for x in 0..input.w() {
-            if input[(x, y)] == b'O' {
-                tot += input.h() - y;
-            }
+    (rocks.into_boxed_slice(), walls.into_boxed_slice())
+}
+
+fn move_north(rocks: &mut [u128], walls: &[u128]) {
+    for i in 1..rocks.len() - 1 {
+        let mut curr_rockl = take(&mut rocks[i]);
+        let mut j = i;
+        while curr_rockl != 0 {
+            let stuck_rockl = curr_rockl & (walls[j - 1] | rocks[j - 1]);
+            rocks[j] |= stuck_rockl;
+            curr_rockl ^= stuck_rockl;
+            j -= 1;
         }
     }
-    tot
 }
 
-pub fn part2(input: &Input) -> usize {
-    let mut seen = FxHashMap::default();
-
-    let mut last_acc = vec![0; input.w()];
-    let mut input = input.clone();
-
-    let mut curr = 0;
-    while curr < 1000000000 {
-        seen.insert(input.clone(), curr);
-
-        last_acc.fill(0);
-        for y in 0..input.h() {
-            for (x, last) in last_acc.iter_mut().enumerate() {
-                if input[(x, y)] == b'O' {
-                    input[(x, y)] = b'.';
-                    input[(x, *last)] = b'O';
-                    *last += 1;
-                } else if input[(x, y)] == b'#' {
-                    *last = y + 1;
-                }
-            }
+fn move_south(rocks: &mut [u128], walls: &[u128]) {
+    for i in (1..rocks.len() - 1).rev() {
+        let mut curr_rockl = take(&mut rocks[i]);
+        let mut j = i;
+        while curr_rockl != 0 {
+            let stuck_rockl = curr_rockl & (walls[j + 1] | rocks[j + 1]);
+            rocks[j] |= stuck_rockl;
+            curr_rockl ^= stuck_rockl;
+            j += 1;
         }
+    }
+}
 
-        for y in 0..input.h() {
-            let mut last = 0;
-            for x in 0..input.w() {
-                if input[(x, y)] == b'O' {
-                    input[(x, y)] = b'.';
-                    input[(last, y)] = b'O';
-                    last += 1;
-                } else if input[(x, y)] == b'#' {
-                    last = x + 1;
-                }
+fn move_west(rocks: &mut [u128], walls: &[u128]) {
+    for i in 1..rocks.len() - 1 {
+        let (mut curr_rockl, mut new_rockl) = (rocks[i], 0);
+        let mut curr_stuck = walls[i];
+        while curr_rockl != 0 {
+            let mut new_stuck = curr_rockl & (curr_stuck >> 1);
+            while new_stuck != 0 {
+                curr_stuck |= new_stuck;
+                new_stuck = (curr_rockl & (curr_stuck >> 1)) & !curr_stuck;
             }
+            new_rockl |= curr_stuck & !walls[i];
+            curr_rockl = (curr_rockl & !new_rockl) << 1;
         }
+        rocks[i] = new_rockl;
+    }
+}
 
-        last_acc.fill(input.h());
-        for y in (0..input.h()).rev() {
-            for (x, last) in last_acc.iter_mut().enumerate().rev() {
-                if input[(x, y)] == b'O' {
-                    *last -= 1;
-                    input[(x, y)] = b'.';
-                    input[(x, *last)] = b'O';
-                } else if input[(x, y)] == b'#' {
-                    *last = y;
-                }
+fn move_east(rocks: &mut [u128], walls: &[u128]) {
+    for i in 1..rocks.len() - 1 {
+        let (mut curr_rockl, mut new_rockl) = (rocks[i], 0);
+        let mut curr_stuck = walls[i];
+        while curr_rockl != 0 {
+            let mut new_stuck = curr_rockl & (curr_stuck << 1);
+            while new_stuck != 0 {
+                curr_stuck |= new_stuck;
+                new_stuck = (curr_rockl & (curr_stuck << 1)) & !curr_stuck;
             }
+            new_rockl |= curr_stuck & !walls[i];
+            curr_rockl = (curr_rockl & !new_rockl) >> 1;
         }
+        rocks[i] = new_rockl;
+    }
+}
 
-        for y in (0..input.h()).rev() {
-            let mut last = input.w();
-            for x in (0..input.w()).rev() {
-                if input[(x, y)] == b'O' {
-                    last -= 1;
-                    input[(x, y)] = b'.';
-                    input[(last, y)] = b'O';
-                } else if input[(x, y)] == b'#' {
-                    last = x;
-                }
-            }
-        }
+fn load_of(rocks: &[u128]) -> usize {
+    rocks
+        .iter()
+        .enumerate()
+        .map(|(i, rockl)| rockl.count_ones() as usize * (rocks.len() - i - 1))
+        .sum()
+}
+
+pub fn part1(input: &Input) -> usize {
+    let (rocks, walls) = input;
+    let mut rocks = rocks.clone();
+    move_north(&mut rocks, walls);
+    load_of(&rocks)
+}
+
+pub fn part2(input: &Input) -> usize {
+    let (rocks, walls) = input;
+    let mut rocks = rocks.clone();
+
+    let mut seen = FxIndexMap::default();
+    let mut curr = 0;
+    while curr < 1000000000 {
+        seen.insert(rocks.clone(), curr);
+
+        move_north(&mut rocks, walls);
+        move_west(&mut rocks, walls);
+        move_south(&mut rocks, walls);
+        move_east(&mut rocks, walls);
 
         curr += 1;
 
-        if let Some(&old) = seen.get(&input) {
+        if let Some(&old) = seen.get(&rocks) {
             let cycle_len = curr - old;
             let cycle_offset = old;
             let answer = (1000000000 - cycle_offset) % cycle_len + cycle_offset;
-            input = seen
-                .drain()
+
+            rocks = seen
+                .drain(..)
                 .find(|&(_, i)| i == answer)
-                .map(|(grid, _)| grid)
+                .map(|(rocks, _)| rocks)
                 .unwrap();
             break;
         }
     }
 
-    let mut tot = 0;
-    for y in 0..input.h() {
-        for x in 0..input.w() {
-            if input[(x, y)] == b'O' {
-                tot += input.h() - y;
-            }
-        }
-    }
-    tot
+    load_of(&rocks)
 }