From 1f5ace53ed4dccfbff2847db3e26acc560d91a3b Mon Sep 17 00:00:00 2001
From: Connor Slade <connor@connorcode.com>
Date: Fri, 22 Dec 2023 01:16:46 -0500
Subject: [PATCH] [2023] Complete day 22

---
 aoc_2023/src/day_22.rs | 63 +++++++++++++++++++++++++++---------------
 1 file changed, 40 insertions(+), 23 deletions(-)

diff --git a/aoc_2023/src/day_22.rs b/aoc_2023/src/day_22.rs
index 2b2cd71..7d8d13b 100644
--- a/aoc_2023/src/day_22.rs
+++ b/aoc_2023/src/day_22.rs
@@ -1,8 +1,11 @@
+use std::collections::HashSet;
+
 use common::{Answer, Solution};
 use nd_vec::{vector, Vec3};
 
 pub struct Day22;
 
+// This could be (and will be in the future) updated to make use of the fact that elements can be sorted by their z value
 impl Solution for Day22 {
     fn name(&self) -> &'static str {
         "Sand Slabs"
@@ -10,14 +13,13 @@ impl Solution for Day22 {
 
     fn part_a(&self, input: &str) -> Answer {
         let mut map = parse(input);
-        map.sort_by_key(|b| b.a.z());
-        while shift_down(&mut map) {}
+        while shift_down(&mut map, false) != 0 {}
 
         let mut count = 0;
         for i in 0..map.len() {
             let mut map_clone = map.clone();
             map_clone.remove(i);
-            if !shift_down(&mut map_clone) {
+            if shift_down(&mut map_clone, false) == 0 {
                 count += 1;
             }
         }
@@ -26,32 +28,47 @@ impl Solution for Day22 {
     }
 
     fn part_b(&self, input: &str) -> Answer {
-        Answer::Unimplemented
+        let mut map = parse(input);
+        while shift_down(&mut map, false) != 0 {}
+
+        let mut count = 0;
+        for i in 0..map.len() {
+            let mut map_clone = map.clone();
+            map_clone.remove(i);
+            count += shift_down(&mut map_clone, true);
+        }
+
+        count.into()
     }
 }
 
-fn shift_down(map: &mut Vec<Box>) -> bool {
-    let mut dirty = false;
-    'outer: for i in 0..map.len() {
-        // If no other box below this one, move it down
-        let item = map[i];
-        for x in item.a.x()..=item.b.x() {
-            for y in item.a.y()..=item.b.y() {
-                if item.a.z() == 1
-                    || map
-                        .iter()
-                        .any(|b| b.contains(vector!(x, y, item.a.z() - 1)))
-                {
-                    continue 'outer;
+fn shift_down(map: &mut Vec<Box>, exhaustive: bool) -> u32 {
+    let mut moved = HashSet::new();
+    let mut dirty = true;
+    while dirty {
+        dirty = false;
+        'outer: for i in 0..map.len() {
+            // If no other box below this one, move it down
+            let item = map[i];
+            for x in item.a.x()..=item.b.x() {
+                for y in item.a.y()..=item.b.y() {
+                    if item.a.z() == 1
+                        || map
+                            .iter()
+                            .any(|b| b.contains(vector!(x, y, item.a.z() - 1)))
+                    {
+                        continue 'outer;
+                    }
                 }
             }
-        }
 
-        map[i].a -= vector!(0, 0, 1);
-        map[i].b -= vector!(0, 0, 1);
-        dirty = true;
+            map[i].a -= vector!(0, 0, 1);
+            map[i].b -= vector!(0, 0, 1);
+            moved.insert(i);
+            dirty = exhaustive;
+        }
     }
-    dirty
+    moved.len() as u32
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -116,6 +133,6 @@ mod test {
 
     #[test]
     fn part_b() {
-        assert_eq!(Day22.part_b(CASE), ().into());
+        assert_eq!(Day22.part_b(CASE), 7.into());
     }
 }