AoC 2024 day 15 part 2

aoc2024/src/day15/python/solution.py

@@ -22,6 +22,7 @@ class BaseWarehouse:
   _ROBOT = '@'
   _WALL = '#'
   _EMPTY = '.'
+  _BOX = 'O'
   _DIR_TO_OPPOSITE = {
     (0, 1): (0, -1),  # Right to left.
     (0, -1): (0, 1),  # Left to right.
@@ -40,12 +41,27 @@ def __init__(self, warehouse_map: WarehouseMap, robot_moves: Sequence[Position])
   def simulate(self) -> None:
     """Simulates all robot moves in the warehouse."""
     while self._next_move_index < len(self._robot_moves):
+      # self._print_map()
       self._move_robot()
       self._next_move_index += 1
 
   def sum_all_boxes_gps_coordinates(self) -> int:
     raise NotImplementedError()
 
+  def _print_map(self) -> None:
+    """Prints the warehouse map to stdout."""
+    dir_map = {
+      (-1, 0): '<',
+      (1, 0): '>',
+      (0, -1): '^',
+      (0, 1): 'v',
+    }
+    for y in range(self._height):
+      print(''.join(self._warehouse_map[y]))
+    next_move = self._robot_moves[self._next_move_index]
+    print(f'Move {dir_map.get(next_move)}')
+    print('')
+
   def _find_robot_pos(self) -> Position:
     """Returns the position of the robot."""
     for y in range(self._height):
@@ -56,9 +72,20 @@ def _find_robot_pos(self) -> Position:
 
   def _move_robot(self) -> None:
     """Attempts to move robot one step. If the robot cannot move, it is a NOP."""
-    raise NotImplementedError()
+    next_move = self._robot_moves[self._next_move_index]
+    next_pos = (self._robot_pos[0] + next_move[0], self._robot_pos[1] + next_move[1])
+    if self._charAt(next_pos) == self._EMPTY:
+      self._update(self._robot_pos, self._EMPTY)
+      self._robot_pos = next_pos
+      self._update(self._robot_pos, self._ROBOT)
+    elif self._charAt(next_pos) in self._BOX:
+      self._push(box_pos=next_pos, dir=next_move)
+      if self._charAt(next_pos) == self._EMPTY:
+        self._update(self._robot_pos, self._EMPTY)
+        self._robot_pos = next_pos
+        self._update(self._robot_pos, self._ROBOT)
 
-  def _push(self, pos: Position, dir: Direction) -> None:
+  def _push(self, box_pos: Position, dir: Direction) -> None:
     raise NotImplementedError()
 
   def _charAt(self, pos: Position) -> str:
@@ -75,8 +102,6 @@ def _within_bounds(self, pos: Position):
 
 
 class Warehouse(BaseWarehouse):
-  _BOX = 'O'
-
   @override
   def sum_all_boxes_gps_coordinates(self) -> int:
     """Returns the sum of all the boxes' GPS coordinates.
@@ -91,33 +116,17 @@ def sum_all_boxes_gps_coordinates(self) -> int:
     return result
 
   @override
-  def _move_robot(self) -> None:
-    """Attempts to move robot one step. If the robot cannot move, it is a NOP."""
-    next_move = self._robot_moves[self._next_move_index]
-    next_pos = (self._robot_pos[0] + next_move[0], self._robot_pos[1] + next_move[1])
-    if self._charAt(next_pos) == self._EMPTY:
-      self._update(self._robot_pos, self._EMPTY)
-      self._robot_pos = next_pos
-      self._update(self._robot_pos, self._ROBOT)
-    elif self._charAt(next_pos) == self._BOX:
-      self._push(pos=next_pos, dir=next_move)
-      if self._charAt(next_pos) == self._EMPTY:
-        self._update(self._robot_pos, self._EMPTY)
-        self._robot_pos = next_pos
-        self._update(self._robot_pos, self._ROBOT)
-
-  @override
-  def _push(self, pos: Position, dir: Direction) -> None:
+  def _push(self, box_pos: Position, dir: Direction) -> None:
     """Attempts to push boxes from given position towards given direction.
     If there is no space to move boxes, it is a NOP."""
-    next_empty_pos = self._find_next_empty_pos(pos, dir)
+    next_empty_pos = self._find_next_empty_pos(box_pos, dir)
     if next_empty_pos is not None:
       next_pos = next_empty_pos
-      while next_pos != pos:
+      while next_pos != box_pos:
         self._update(next_pos, self._BOX)
         opposite_dir = self._DIR_TO_OPPOSITE.get(dir)
         next_pos = (next_pos[0] + opposite_dir[0], next_pos[1] + opposite_dir[1])
-      self._update(pos, self._EMPTY)
+      self._update(box_pos, self._EMPTY)
 
   def _find_next_empty_pos(self, pos: Position, dir: Direction) -> Position | None:
     """Returns the next empty space from a position towards a given direction."""
@@ -131,6 +140,148 @@ def _find_next_empty_pos(self, pos: Position, dir: Direction) -> Position | None
     return None
 
 
+class DoubleWarehouse(BaseWarehouse):
+  _BOX = '[]'
+
+  def __init__(self, warehouse_map: WarehouseMap, robot_moves: Sequence[Position]):
+    super().__init__(warehouse_map, robot_moves)
+    self._resize_warehouse()
+    self._width = len(self._warehouse_map[0])
+    self._height = len(self._warehouse_map)
+    self._robot_pos = self._find_robot_pos()
+    self._robot_moves = robot_moves
+    self._next_move_index = 0
+
+  @override
+  def sum_all_boxes_gps_coordinates(self) -> int:
+    """Returns the sum of all the boxes' GPS coordinates.
+    A box GPS coordinate is 100 times its distance from the top edge of the
+    warehouse plus its distance from the left edge of the map.
+    """
+    result = 0
+    for y in range(self._height):
+      for x in range(self._width):
+        if self._charAt(pos=(x, y)) == self._BOX[0]:
+          result += 100 * y + x
+    return result
+
+  def _resize_warehouse(self):
+    # Resize (double) the map.
+    resized_map = []
+    for y in range(self._height):
+      line = []
+      for x in range(self._width):
+        value = self._charAt(pos=(x, y))
+        resized_value = None
+        if value == self._WALL:
+          resized_value = self._WALL * 2
+        if value == super()._BOX:
+          resized_value = self._BOX
+        if value == self._EMPTY:
+          resized_value = self._EMPTY * 2
+        if value == self._ROBOT:
+          resized_value = self._ROBOT + self._EMPTY
+        line.append(resized_value)
+      resized_map.append(list(''.join(line)))
+    self._warehouse_map = resized_map
+
+  @override
+  def _push(self, box_pos: Position, dir: Direction) -> None:
+    if dir in ((-1, 0), (1, 0)):  # Boxes are 1-tile when moving horizontally.
+      next_empty_pos = self._find_next_empty_x(box_pos, dir)
+      if next_empty_pos is not None:
+        line = ''.join(self._warehouse_map[box_pos[1]])
+        if dir == (-1, 0):  # Shift box(es) left.
+          shifted_left_line = (line[:next_empty_pos[0]]
+                               + line[next_empty_pos[0] + 1:box_pos[0] + 1]
+                               + self._EMPTY
+                               + line[box_pos[0] + 1:])
+          self._warehouse_map[box_pos[1]] = list(shifted_left_line)
+        elif dir == (1, 0):  # Shift box(es) right
+          shifted_right_line = (line[:box_pos[0]]
+                                + self._EMPTY
+                                + line[box_pos[0]:next_empty_pos[0]]
+                                + line[next_empty_pos[0] + 1:])
+          self._warehouse_map[box_pos[1]] = list(shifted_right_line)
+    elif dir in ((0, -1), (0, 1)):  # Boxes are 2-tile when moving vertically.
+      box_side = self._charAt(box_pos)
+      if box_side == '[':
+        if self._can_move_boxes_vertically(left_pos=box_pos, dir=dir):
+          self._move_boxes_vertically(left_pos=box_pos, dir=dir)
+      elif box_side == ']':
+        left_pos = (box_pos[0] - 1, box_pos[1])
+        if self._can_move_boxes_vertically(left_pos, dir):
+          self._move_boxes_vertically(left_pos, dir)
+      else:
+        raise ValueError(f'pos={box_pos} not a box, was: {box_side}')
+    else:
+      raise ValueError(f'Unrecognized dir={dir}')
+
+  def _can_move_boxes_vertically(self, left_pos: Position, dir: Direction) -> bool:
+    """Returns true if all boxes can be moved given a box position."""
+    behind_left_pos = left_pos[0] + dir[0], left_pos[1] + dir[1]
+    behind_right_pos = (behind_left_pos[0] + 1, behind_left_pos[1])
+    behind_value = self._charAt(behind_left_pos) + self._charAt(behind_right_pos)
+    # Base case. If the last adjacent box is against a wall, no box moves.
+    if behind_value == self._WALL * 2:
+      return False
+    if behind_value == self._EMPTY * 2:
+      return True
+    if behind_value == self._BOX:
+      return self._can_move_boxes_vertically(left_pos=behind_left_pos, dir=dir)
+    if behind_value == self._EMPTY + self._BOX[0]:
+      # Touching left side of another box.
+      return self._can_move_boxes_vertically(left_pos=(behind_left_pos[0] + 1, behind_left_pos[1]), dir=dir)
+    if behind_value == self._BOX[1] + self._EMPTY:
+      # Touching right side of another box.
+      return self._can_move_boxes_vertically(left_pos=(behind_left_pos[0] - 1, behind_left_pos[1]), dir=dir)
+    if behind_value == self._BOX[::-1]:  # 2 boxes behind!
+      return (self._can_move_boxes_vertically(left_pos=(behind_left_pos[0] - 1, behind_left_pos[1]), dir=dir)
+              and self._can_move_boxes_vertically(left_pos=(behind_left_pos[0] + 1, behind_left_pos[1]), dir=dir))
+
+  def _move_boxes_vertically(self, left_pos: Position, dir: Direction) -> None:
+    """Moves all boxes that can be moved in the y-axis given direction."""
+    behind_left_pos = left_pos[0] + dir[0], left_pos[1] + dir[1]
+    behind_right_pos = (behind_left_pos[0] + 1, behind_left_pos[1])
+    behind_value = self._charAt(behind_left_pos) + self._charAt(behind_right_pos)
+    # Base case.
+    if behind_value == self._EMPTY * 2:
+      self._move_box_vertically(left_pos, dir)
+      return
+    if behind_value == self._BOX:
+      self._move_boxes_vertically(left_pos=behind_left_pos, dir=dir)
+      self._move_box_vertically(left_pos, dir)
+    if behind_value == self._EMPTY + self._BOX[0]:
+      self._move_boxes_vertically(left_pos=(behind_left_pos[0] + 1, behind_left_pos[1]), dir=dir)
+      self._move_box_vertically(left_pos, dir)
+    if behind_value == self._BOX[1] + self._EMPTY:
+      self._move_boxes_vertically(left_pos=(behind_left_pos[0] - 1, behind_left_pos[1]), dir=dir)
+      self._move_box_vertically(left_pos, dir)
+    if behind_value == self._BOX[::-1]:  # 2 boxes behind!
+      self._move_boxes_vertically(left_pos=(behind_left_pos[0] - 1, behind_left_pos[1]), dir=dir)
+      self._move_boxes_vertically(left_pos=(behind_left_pos[0] + 1, behind_left_pos[1]), dir=dir)
+      self._move_box_vertically(left_pos, dir)
+
+  def _move_box_vertically(self, left_pos: Position, dir: Direction) -> None:
+    """Moves box given up or down."""
+    behind_left_pos = left_pos[0] + dir[0], left_pos[1] + dir[1]
+    behind_right_pos = (behind_left_pos[0] + 1, behind_left_pos[1])
+    self._update(pos=behind_left_pos, value=self._BOX[0])
+    self._update(pos=behind_right_pos, value=self._BOX[1])
+    self._update(pos=left_pos, value=self._EMPTY)
+    self._update(pos=(left_pos[0] + 1, left_pos[1]), value=self._EMPTY)
+
+  def _find_next_empty_x(self, pos: Position, dir: Direction) -> Position | None:
+    while self._within_bounds(pos):
+      if self._charAt(pos) == self._WALL:
+        return None  # We hit a wall before an empty space.
+      if self._charAt(pos) == self._EMPTY:
+        return pos
+      pos = (pos[0] + dir[0], pos[1] + dir[1])
+    # Empty space was not found.
+    return None
+
+
 def _parse(input: Sequence[str]) -> tuple[WarehouseMap, Sequence[Position]]:
   """Parses the input and returns the warehouse map and sequence of robot moves."""
   is_map = True
@@ -154,10 +305,10 @@ def _parse(input: Sequence[str]) -> tuple[WarehouseMap, Sequence[Position]]:
   return tuple(warehouse_map), tuple(robot_moves)
 
 
-def sum_all_boxes_gps_coordinates(input: Sequence[str]) -> int:
-  """Simulates all the robot moves in the warehouse and returns the sum of
-  all boxes GPS coordinates at the end state."""
+def sum_all_boxes_gps_coordinates(input: Sequence[str], double_size=False) -> int:
+  """Simulates all the robot moves in the warehouse and
+  returns the sum of all boxes GPS coordinates at the end state."""
   warehouse_map, robot_moves = _parse(input)
-  warehouse = Warehouse(warehouse_map, robot_moves)
+  warehouse = DoubleWarehouse(warehouse_map, robot_moves) if double_size else Warehouse(warehouse_map, robot_moves)
   warehouse.simulate()
   return warehouse.sum_all_boxes_gps_coordinates()

aoc2024/test/day15/python/example3.txt

@@ -0,0 +1,9 @@
+#######
+#...#.#
+#.....#
+#..OO@#
+#..O..#
+#.....#
+#######
+
+<vv<<^^<<^^

aoc2024/test/day15/python/test_solution.py

@@ -31,6 +31,15 @@ def test_part1_withLargerExample_correct(self):
   def test_part1_withPuzzleInput_correct(self):
     self.assertEqual(1383666, sum_all_boxes_gps_coordinates(self.input))
 
+  def test_part2_withAnotherExample_correct(self):
+    self.assertEqual(618, sum_all_boxes_gps_coordinates(self.examples[2], double_size=True))
+
+  def test_part2_withLargerExample_correct(self):
+    self.assertEqual(9021, sum_all_boxes_gps_coordinates(self.examples[1], double_size=True))
+
+  def test_part2_withPuzzleInput_correct(self):
+    self.assertEqual(1412866, sum_all_boxes_gps_coordinates(self.input, double_size=True))
+
 
 if __name__ == '__main__':
   unittest.main()