Day 18 part 2. This needed some optimization

1 files changed, 244 insertions, 71 deletions

diff --git a/src/bin/day_18.rs b/src/bin/day_18.rs
index 71886a2..8714e4f 100644
--- a/src/bin/day_18.rs
+++ b/src/bin/day_18.rs
@@ -1,6 +1,8 @@
+use rayon::prelude::*;
 use rpds::rbt_set;
 use rpds::vector;
 use rpds::vector::Vector;
+use rpds::RedBlackTreeMap;
 use rpds::RedBlackTreeSet;
 use std::io;
 use std::io::prelude::*;
@@ -18,7 +20,7 @@ struct Opt {}
 
 fn main() {
     let stdin = io::stdin();
-    let opt = Opt::from_args();
+    let _opt = Opt::from_args();
 
     let maze: Maze = stdin
         .lock()
@@ -41,8 +43,8 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
 
 struct Maze {
     blocks: Vec<Vec<char>>,
-    start: Location,
-    keys: RedBlackTreeSet<char>,
+    start: BoardState,
+    keys: KeySet,
 }
 
 impl FromIterator<String> for Maze {
@@ -58,27 +60,23 @@ impl FromIterator<String> for Maze {
 impl Maze {
     fn from_blocks(blocks: Vec<Vec<char>>) -> Maze {
         Maze {
-            start: blocks
-                .iter()
-                .enumerate()
-                .filter_map(|(y, line)| {
-                    line.iter()
-                        .enumerate()
-                        .find(|(x, ch)| **ch == '@')
-                        .map(|(x, _ch)| Location {
-                            x,
-                            y,
-                            keys: RedBlackTreeSet::new(),
-                        })
-                })
-                .next()
-                .unwrap_or(Location::default()),
+            start: BoardState {
+                robots: blocks
+                    .iter()
+                    .enumerate()
+                    .flat_map(|(y, line)| {
+                        line.iter()
+                            .enumerate()
+                            .filter(|(_x, ch)| **ch == '@')
+                            .map(move |(x, _ch)| Location { x, y })
+                    })
+                    .collect(),
+                keys: KeySet::default(),
+            },
             keys: blocks
                 .iter()
                 .flat_map(|line| line.iter())
-                .filter(|c| c.is_lowercase() && c.is_alphabetic())
-                .cloned()
-                .collect(),
+                .fold(KeySet::default(), |acc, next| acc.add_key(*next)),
             blocks,
         }
     }
@@ -92,7 +90,10 @@ impl Maze {
 
     fn shortest_route(&self) -> usize {
         iter::successors(
-            Some((rbt_set![self.start.clone()], vector![self.start.clone()])),
+            Some((
+                ExploredStates::default().insert(&self.start),
+                rbt_set![self.start.clone()],
+            )),
             |(explored, locations)| {
                 Some(Maze::next_depth_states(
                     explored.clone(),
@@ -100,69 +101,138 @@ impl Maze {
                 ))
             },
         )
-        // .inspect(|(explored, locations)| eprintln!("{:?}", explored.size()))
-        .take_while(|(_explored, locations)| locations.len() > 0)
+        // .inspect(|(explored, states)| eprintln!("{:?}", states))
+        .take_while(|(_explored, states)| states.size() > 0)
         .enumerate()
-        .find(|(_i, (_explored, locations))| {
-            locations.iter().any(|location| location.keys == self.keys)
-        })
+        .find(|(_i, (_explored, states))| states.iter().any(|state| state.keys == self.keys))
         .unwrap()
         .0
     }
 
     fn next_depth_states(
-        explored: RedBlackTreeSet<Location>,
-        locations: Vector<Location>,
-    ) -> (RedBlackTreeSet<Location>, Vector<Location>) {
+        explored: ExploredStates,
+        locations: RedBlackTreeSet<BoardState>,
+    ) -> (ExploredStates, RedBlackTreeSet<BoardState>) {
         (
             locations
                 .iter()
-                .fold(explored, |acc, next| acc.insert(next.clone())),
-            Maze::dedup_locations(locations),
+                .fold(explored, |acc, next| acc.insert(next)),
+            locations,
         )
     }
 
-    fn dedup_locations(locations: Vector<Location>) -> Vector<Location> {
-        locations.iter().fold(Vector::new(), |acc, next| {
-            if acc.iter().any(|loc| loc == next) {
-                acc
-            } else {
-                acc.push_back(next.clone())
-            }
-        })
-    }
-
     fn next_locations(
         &self,
-        locations: &Vector<Location>,
-        explored: &RedBlackTreeSet<Location>,
-    ) -> Vector<Location> {
+        locations: &RedBlackTreeSet<BoardState>,
+        explored: &ExploredStates,
+    ) -> RedBlackTreeSet<BoardState> {
         locations
             .iter()
             .flat_map(|current| {
-                [(-1, 0), (1, 0), (0, -1), (0, 1)]
-                    .iter()
-                    .map(move |(dx, dy)| current.next(*dx, *dy))
-                    .map(move |location| {
-                        self.block_at(&location)
-                            .map(|c| location.with_key(c))
-                            .unwrap_or(location)
-                    })
+                (0..current.robots.len()).flat_map(move |i_robot| {
+                    [(-1, 0), (1, 0), (0, -1), (0, 1)]
+                        .iter()
+                        .map(move |(dx, dy)| current.next(i_robot, *dx, *dy))
+                        .map(move |(state, new_location)| {
+                            self.block_at(&new_location)
+                                .map(|c| state.with_key(c))
+                                .unwrap_or(state)
+                        })
+                })
             })
-            .filter(|location| self.is_open(location))
-            .filter(|location| !explored.contains(location))
+            .filter(|state| self.is_open(state))
+            .filter(|state| !explored.contains(state))
             .collect()
     }
 
-    fn is_open(&self, location: &Location) -> bool {
-        self.block_at(location)
-            .map(|c| Maze::char_is_open(c, &location.keys))
-            .unwrap_or(false)
+    fn is_open(&self, state: &BoardState) -> bool {
+        state.robots.iter().all(|location| {
+            self.block_at(location)
+                .map(|c| Maze::char_is_open(c, state.keys))
+                .unwrap_or(false)
+        })
+    }
+
+    fn char_is_open(c: char, keys: KeySet) -> bool {
+        c != '#' && !keys.door_is_locked(c)
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+struct ExploredStates {
+    for_key: RedBlackTreeMap<KeySet, ExploredStatesForKey>,
+}
+
+#[derive(Default, Debug, Clone)]
+struct ExploredStatesForKey {
+    robots: Vector<RedBlackTreeSet<Location>>,
+}
+
+impl ExploredStates {
+    fn insert(&self, new: &BoardState) -> ExploredStates {
+        ExploredStates {
+            for_key: self.for_key.insert(
+                new.keys.clone(),
+                self.for_key
+                    .get(&new.keys)
+                    .map(|current| ExploredStatesForKey {
+                        robots: current
+                            .robots
+                            .iter()
+                            .zip(new.robots.iter())
+                            .map(|(current_explored, robot)| current_explored.insert(robot.clone()))
+                            .collect(),
+                    })
+                    .unwrap_or_else(|| ExploredStatesForKey {
+                        robots: new
+                            .robots
+                            .iter()
+                            .map(|robot| RedBlackTreeSet::new().insert(robot.clone()))
+                            .collect(),
+                    }),
+            ),
+        }
+    }
+
+    fn contains(&self, state: &BoardState) -> bool {
+        self.for_key
+            .get(&state.keys)
+            .filter(|for_key| {
+                state
+                    .robots
+                    .iter()
+                    .zip(for_key.robots.iter())
+                    .all(|(r_state, r_explored)| r_explored.contains(r_state))
+            })
+            .is_some()
+    }
+}
+
+#[derive(Default, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+struct BoardState {
+    robots: Vector<Location>,
+    keys: KeySet,
+}
+
+impl BoardState {
+    fn next(&self, i_robot: usize, dx: isize, dy: isize) -> (BoardState, Location) {
+        (
+            BoardState {
+                robots: self
+                    .robots
+                    .set(i_robot, self.robots[i_robot].next(dx, dy))
+                    .unwrap(),
+                ..self.clone()
+            },
+            self.robots[i_robot].next(dx, dy),
+        )
     }
 
-    fn char_is_open(c: char, keys: &RedBlackTreeSet<char>) -> bool {
-        c != '#'
-            && (!c.is_uppercase() || !c.is_alphabetic() || keys.contains(&c.to_ascii_lowercase()))
+    fn with_key(&self, c: char) -> BoardState {
+        BoardState {
+            keys: self.keys.add_key(c),
+            ..self.clone()
+        }
     }
 }
 
@@ -170,7 +240,6 @@ impl Maze {
 struct Location {
     x: usize,
     y: usize,
-    keys: RedBlackTreeSet<char>,
 }
 
 impl Location {
@@ -178,17 +247,121 @@ impl Location {
         Location {
             x: (self.x as isize + dx) as usize,
             y: (self.y as isize + dy) as usize,
-            ..self.clone()
         }
     }
-    fn with_key(&self, c: char) -> Location {
-        if c.is_lowercase() && c.is_alphabetic() {
-            Location {
-                keys: self.keys.insert(c),
-                ..self.clone()
-            }
+}
+
+#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+struct KeySet {
+    bitset: u32,
+}
+
+impl KeySet {
+    fn add_key(self, key: char) -> KeySet {
+        KeySet {
+            bitset: self.bitset | KeySet::key_to_bitset(key),
+        }
+    }
+
+    fn door_is_locked(self, door: char) -> bool {
+        KeySet::door_to_bitset(door) & (!self.bitset) > 0
+    }
+
+    fn key_to_bitset(key: char) -> u32 {
+        if key.is_ascii_alphabetic() && key.is_ascii_lowercase() {
+            1 << (key as u8 - b'a')
         } else {
-            self.clone()
+            0
         }
     }
+
+    fn door_to_bitset(door: char) -> u32 {
+        if door.is_ascii_alphabetic() && door.is_ascii_uppercase() {
+            1 << (door as u8 - b'A')
+        } else {
+            0
+        }
+    }
+}
+
+#[test]
+fn doors_are_locked_without_key() {
+    assert_eq!(true, KeySet::default().door_is_locked('A'))
+}
+
+#[test]
+fn doors_are_unlocked_with_key() {
+    assert_eq!(false, KeySet::default().add_key('a').door_is_locked('A'))
+}
+
+#[test]
+fn example_1() {
+    let maze: Maze = r"
+#########
+#b.A.@.a#
+#########
+"
+    .split('\n')
+    .map(|s| s.to_string())
+    .collect();
+
+    assert_eq!(maze.shortest_route(), 8);
+}
+
+#[test]
+fn example_2() {
+    let maze: Maze = r"
+#################
+#i.G..c...e..H.p#
+########.########
+#j.A..b...f..D.o#
+########@########
+#k.E..a...g..B.n#
+########.########
+#l.F..d...h..C.m#
+#################
+"
+    .split('\n')
+    .map(|s| s.to_string())
+    .collect();
+
+    assert_eq!(maze.shortest_route(), 136);
+}
+
+#[test]
+fn example_3() {
+    let maze: Maze = r"
+#############
+#DcBa.#.GhKl#
+#.###@#@#I###
+#e#d#####j#k#
+###C#@#@###J#
+#fEbA.#.FgHi#
+#############
+"
+    .split('\n')
+    .map(|s| s.to_string())
+    .collect();
+
+    assert_eq!(maze.shortest_route(), 32);
+}
+
+#[test]
+fn example_4() {
+    let maze: Maze = r"
+#############
+#g#f.D#..h#l#
+#F###e#E###.#
+#dCba@#@BcIJ#
+#############
+#nK.L@#@G...#
+#M###N#H###.#
+#o#m..#i#jk.#
+#############
+"
+    .split('\n')
+    .map(|s| s.to_string())
+    .collect();
+
+    assert_eq!(maze.shortest_route(), 74);
 }