Day 20 part 2 - now with depth!

1 files changed, 52 insertions, 70 deletions

diff --git a/src/bin/day_20.rs b/src/bin/day_20.rs
index 1bce235..953df75 100644
--- a/src/bin/day_20.rs
+++ b/src/bin/day_20.rs
@@ -14,11 +14,15 @@ use structopt::StructOpt;
 /// Finds the shortest path through a maze with portals.
 ///
 /// See https://adventofcode.com/2019/day/20 for details.
-struct Opt {}
+struct Opt {
+    /// include the rule that going through portals changes your depth
+    #[structopt(short = "d")]
+    include_depth: bool,
+}
 
 fn main() {
     let stdin = io::stdin();
-    let _opt = Opt::from_args();
+    let opt = Opt::from_args();
 
     let maze = stdin
         .lock()
@@ -27,7 +31,7 @@ fn main() {
         .collect::<MazeBuilder>()
         .build();
 
-    println!("{}", maze.shortest_route());
+    println!("{}", maze.shortest_route(opt.include_depth));
 }
 
 fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
@@ -87,7 +91,11 @@ impl MazeBuilder {
                 ),
                  (next_id, next_p)| match unmatched.get(&next_id) {
                     Some(pair) => (
-                        matched.push_back((pair.clone(), next_p)),
+                        matched.push_back(pair.clone().inside_out(
+                            next_p,
+                            self.map[0].len(),
+                            self.map.len(),
+                        )),
                         unmatched.remove(&next_id),
                     ),
                     None => (matched, unmatched.insert(next_id, next_p)),
@@ -183,34 +191,53 @@ impl Point {
             y: (self.y as isize + y) as usize,
         }
     }
+
+    fn inside_out(self, other: Point, width: usize, height: usize) -> (Point, Point) {
+        if self.closest_side(width, height) > other.closest_side(width, height) {
+            (self, other)
+        } else {
+            (other, self)
+        }
+    }
+
+    fn closest_side(&self, width: usize, height: usize) -> usize {
+        self.x.min(width - self.x).min(self.y).min(height - self.y)
+    }
 }
 
 impl Maze {
-    fn shortest_route(&self) -> usize {
+    fn shortest_route(&self, include_depth: bool) -> usize {
         iter::successors(
             Some((
-                rbt_set![self.entrance.clone()],
-                rbt_set![self.entrance.clone()],
+                rbt_set![(self.entrance.clone(), 0)],
+                rbt_set![(self.entrance.clone(), 0)],
             )),
             |(explored, locations)| {
                 Some(Maze::next_depth_states(
                     explored.clone(),
-                    self.next_locations(locations, explored),
+                    self.next_locations(locations, explored, include_depth),
                 ))
             },
         )
         // .inspect(|(explored, states)| eprintln!("{:?}", states))
         .take_while(|(_explored, states)| states.size() > 0)
         .enumerate()
-        .find(|(_i, (_explored, states))| states.iter().any(|state| *state == self.exit))
+        .find(|(_i, (_explored, states))| {
+            states
+                .iter()
+                .any(|(p_state, depth)| *p_state == self.exit && (!include_depth || *depth == 0))
+        })
         .unwrap()
         .0
     }
 
     fn next_depth_states(
-        explored: RedBlackTreeSet<Point>,
-        locations: RedBlackTreeSet<Point>,
-    ) -> (RedBlackTreeSet<Point>, RedBlackTreeSet<Point>) {
+        explored: RedBlackTreeSet<(Point, isize)>,
+        locations: RedBlackTreeSet<(Point, isize)>,
+    ) -> (
+        RedBlackTreeSet<(Point, isize)>,
+        RedBlackTreeSet<(Point, isize)>,
+    ) {
         (
             locations
                 .iter()
@@ -221,31 +248,32 @@ impl Maze {
 
     fn next_locations(
         &self,
-        locations: &RedBlackTreeSet<Point>,
-        explored: &RedBlackTreeSet<Point>,
-    ) -> RedBlackTreeSet<Point> {
+        locations: &RedBlackTreeSet<(Point, isize)>,
+        explored: &RedBlackTreeSet<(Point, isize)>,
+        include_depth: bool,
+    ) -> RedBlackTreeSet<(Point, isize)> {
         locations
             .iter()
-            .flat_map(|p| {
+            .flat_map(|(p, depth)| {
                 [(-1, 0), (1, 0), (0, -1), (0, 1)]
                     .iter()
-                    .map(move |(dx, dy)| p.add(*dx, *dy))
+                    .map(move |(dx, dy)| (p.add(*dx, *dy), *depth))
                     .chain(
                         self.portals
                             .iter()
                             .filter(move |(from, _to)| p == from)
-                            .map(|(_from, to)| to)
-                            .cloned(),
+                            .map(move |(_from, to)| (to.clone(), depth + 1)),
                     )
                     .chain(
                         self.portals
                             .iter()
                             .filter(move |(_to, from)| p == from)
-                            .map(|(to, _from)| to)
-                            .cloned(),
+                            .map(move |(to, _from)| (to.clone(), depth - 1)),
                     )
             })
-            .filter(|p_next| !self.walls[p_next.y][p_next.x])
+            .filter(|(p_next, depth)| {
+                !self.walls[p_next.y][p_next.x] && (!include_depth || *depth >= 0)
+            })
             .filter(|state| !explored.contains(state))
             .collect()
     }
@@ -277,52 +305,6 @@ FG..#########.....#
         .collect::<MazeBuilder>()
         .build();
 
-    assert_eq!(maze.shortest_route(), 23);
-}
-
-#[test]
-fn portal_maze_example_2() {
-    let maze: Maze = r"                   A               
-                   A               
-  #################.#############  
-  #.#...#...................#.#.#  
-  #.#.#.###.###.###.#########.#.#  
-  #.#.#.......#...#.....#.#.#...#  
-  #.#########.###.#####.#.#.###.#  
-  #.............#.#.....#.......#  
-  ###.###########.###.#####.#.#.#  
-  #.....#        A   C    #.#.#.#  
-  #######        S   P    #####.#  
-  #.#...#                 #......VT
-  #.#.#.#                 #.#####  
-  #...#.#               YN....#.#  
-  #.###.#                 #####.#  
-DI....#.#                 #.....#  
-  #####.#                 #.###.#  
-ZZ......#               QG....#..AS
-  ###.###                 #######  
-JO..#.#.#                 #.....#  
-  #.#.#.#                 ###.#.#  
-  #...#..DI             BU....#..LF
-  #####.#                 #.#####  
-YN......#               VT..#....QG
-  #.###.#                 #.###.#  
-  #.#...#                 #.....#  
-  ###.###    J L     J    #.#.###  
-  #.....#    O F     P    #.#...#  
-  #.###.#####.#.#####.#####.###.#  
-  #...#.#.#...#.....#.....#.#...#  
-  #.#####.###.###.#.#.#########.#  
-  #...#.#.....#...#.#.#.#.....#.#  
-  #.###.#####.###.###.#.#.#######  
-  #.#.........#...#.............#  
-  #########.###.###.#############  
-           B   J   C               
-           U   P   P               "
-        .split('\n')
-        .map(|s| s.to_string())
-        .collect::<MazeBuilder>()
-        .build();
-
-    assert_eq!(maze.shortest_route(), 58);
+    assert_eq!(maze.shortest_route(false), 23);
+    assert_eq!(maze.shortest_route(true), 26);
 }