commit 50189f6b23b0e0553d41079ba54cc3f588cad4d6
parent 2f41fa16e2b71497ab2599939664efe08cff03ee
Author: Eamon Caddigan <eamon.caddigan@gmail.com>
Date: Mon, 27 Dec 2021 12:47:14 -0500
Solution to day 23, part 2
Diffstat:
| A | day23_part2.py | | | 398 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
1 file changed, 398 insertions(+), 0 deletions(-)
diff --git a/day23_part2.py b/day23_part2.py
@@ -0,0 +1,398 @@
+#!/usr/bin/env python
+"""Advent of Code 2021, day 23 (part 2): Amphipod
+Play a shrimp shuffling game"""
+
+# This began as a refactoring of the code that I wrote for part 1. Even before
+# implementing a heuristic function (i.e., when this was essentially running
+# Dijkstra's algorithm with a bit of pruning), this version ran much faster
+# (260 vs 686 ms for the example input, and 1.68 vs. 4.1 s for my puzzle
+# input).
+#
+# After adding a proper heuristic function, solution times dropped to 9.81 ms
+# for the example input and 57 ms for my puzzle input. It takes 493 ms to solve
+# part 2 using my puzzle input.
+
+from typing import Tuple, List, Union, Dict, Deque, NewType, cast
+from heapq import heappush, heappop
+from collections import deque
+from math import inf
+
+from day23_part1 import EXAMPLE_INPUT
+from utils import get_puzzle_input
+
+ROOM_SIZE = 4
+
+GameState = NewType(
+ 'GameState',
+ Tuple[int, int, int, int, int, int, int, int,
+ int, int, int, int, int, int, int, int]
+)
+
+def parse_input(input_string: str) -> GameState:
+ """Given the puzzle input, return the initial state of the game"""
+ # This implementation assumes that all amphipods start in a room; it's
+ # unable to parse a game state with amphipods in the hallway
+ lines = input_string.rstrip('\n').split('\n')
+ lines.insert(3, ' #D#C#B#A#')
+ lines.insert(4, ' #D#B#A#C#')
+ state_list: List[Union[None, int]] = [None] * 16
+ for room_line in range(4):
+ amphipods = [4 * (ord(c) - ord('A')) \
+ for c in lines[2 + room_line].replace('#', '').replace(' ', '')]
+ for room_num, amphipod in enumerate(amphipods):
+ while state_list[amphipod] is not None:
+ amphipod += 1
+ state_list[amphipod] = room_line + 4 * room_num + 11
+ return cast(GameState, tuple(state_list))
+
+# Because of my stubborn decision to represent the game state in an awkward
+# way, I need a bunch of little functions to help me actually make sense of the
+# state of the game.
+
+def valid_hall_locations() -> Tuple[int, ...]:
+ """Return all of the hall locations in which an amphipod can legally stop"""
+ #return tuple(l for l in range(11) if l not in range(2, 9, 2))
+ return (0, 1, 3, 5, 7, 9, 10)
+
+def room_entrance_locations() -> Tuple[Union[int, None], ...]:
+ """Return the location in the hallway of the entrance to the rooms"""
+ #return (None,) + tuple(range(2, 9, 2))
+ return (None, 2, 4, 6, 8)
+
+def everybody_home(game_state: GameState) -> bool:
+ """Quickly(ish) determine if all the amphipods are in their home rooms"""
+ return all(a // ROOM_SIZE == (l - 11) // ROOM_SIZE \
+ for a, l in enumerate(game_state))
+
+def location_to_room_spot(location: int) -> Tuple[int, int]:
+ """Convert an 'absolute location' to a room and spot number (the hallway is
+ room 0)"""
+ if location < 11:
+ return (0, location)
+ return divmod(location - (11 - ROOM_SIZE), ROOM_SIZE)
+
+def room_spot_to_location(room: int, spot: int) -> int:
+ """Given a room and spot number, find the location"""
+ if room == 0:
+ return spot
+ return 11 + (room-1) * ROOM_SIZE + spot
+
+def amphipod_type(amphipod: int) -> int:
+ """There are 4 * `ROOM_SIZE` amphipods, but four types; this is useful
+ because the type ID is also the target room number (which is why this is
+ one-based)"""
+ return amphipod // ROOM_SIZE + 1
+
+def get_energy_per_step(amphipod_home: int) -> int:
+ """Different types of amphipod consume different amounts of energy per
+ step"""
+ return cast(int, (None, 1, 10, 100, 1000)[amphipod_home])
+
+def make_state(game_state: GameState, amphipod: int,
+ new_location: int) -> GameState:
+ """Return a new state where the selected amphipod is in the given
+ location"""
+ game_state_list = list(game_state)
+ game_state_list[amphipod] = new_location
+ return cast(GameState, tuple(game_state_list))
+
+# Below are some functions (and their helpers) that return lists of potential
+# moves. A "move" here means a location/number of steps pair (represented as a
+# tuple). All of these return lists even though only `room_to_hallway` has the
+# ability to return more than one move.
+
+def stuck_in_room(room: int, spot: int, location_empty: List[bool]) -> bool:
+ """Helper function to indicate whether a given spot is "stuck" in a room"""
+ locations = (room_spot_to_location(room, s) for s in range(spot-1, -1, -1))
+ # note to self: any([]) returns false and all([]) returns True
+ return not all(location_empty[l] for l in locations)
+
+def path_clear(from_hall_spot: int, to_hall_spot: int,
+ location_empty: List[bool]) -> bool:
+ """Indicates whether the pathway is clear between the originating spot and
+ the destination spot (the originating spot need not be clear)"""
+ if from_hall_spot < to_hall_spot:
+ test_range = range(from_hall_spot + 1, to_hall_spot + 1)
+ else:
+ test_range = range(to_hall_spot, from_hall_spot)
+ return all(location_empty[l] for l in test_range)
+
+def free_room_spot(room: int, location_empty: List[bool]) -> Union[int, None]:
+ """In a room that's ready, return the best (lowest) spot"""
+ for spot in range(ROOM_SIZE-1, -1, -1):
+ if location_empty[room_spot_to_location(room, spot)]:
+ return spot
+ return None
+
+def room_to_hallway(room: int, spot: int,
+ location_empty: List[bool]) -> List[Tuple[int, int]]:
+ """List all the loecal moves from the given room/spot into the hallway"""
+ if stuck_in_room(room, spot, location_empty):
+ return []
+ from_entrance = room_entrance_locations()[room]
+ assert from_entrance is not None
+ # This is extremely hacky, but we'll loop through the list of hall
+ # locations once in each direction.
+ destination_locations = []
+ for location in valid_hall_locations():
+ if location > from_entrance:
+ if not location_empty[location]:
+ break
+ destination_locations.append(location)
+ for location in reversed(valid_hall_locations()):
+ if location < from_entrance:
+ if not location_empty[location]:
+ break
+ destination_locations.append(location)
+ return [(l, abs(from_entrance - l) + spot + 1) \
+ for l in destination_locations]
+
+def hallway_to_room(hall_spot: int, room: int,
+ location_empty: List[bool]) -> List[Tuple[int, int]]:
+ """Return a list containing a move from the given spot in the hallway into
+ the given room, or an empty list if there's no such move"""
+ to_entrance = room_entrance_locations()[room]
+ assert to_entrance is not None
+ if not path_clear(hall_spot, to_entrance, location_empty):
+ return []
+ to_spot = free_room_spot(room, location_empty)
+ assert to_spot is not None
+ return [(room_spot_to_location(room, to_spot),
+ to_spot + abs(hall_spot - to_entrance) + 1)]
+
+def room_to_room(from_room: int, spot: int, to_room: int,
+ location_empty: List[bool]) -> List[Tuple[int, int]]:
+ """Return a list containing a move from the given spot and room into
+ the target room, or an empty list if there's no such move"""
+ from_entrance = room_entrance_locations()[from_room]
+ to_entrance = room_entrance_locations()[to_room]
+ assert from_entrance is not None and to_entrance is not None
+ if stuck_in_room(from_room, spot, location_empty) \
+ or not path_clear(from_entrance, to_entrance, location_empty):
+ return []
+ to_spot = free_room_spot(to_room, location_empty)
+ assert to_spot is not None
+ return [(room_spot_to_location(to_room, to_spot),
+ spot + to_spot + abs(from_entrance - to_entrance) + 2)]
+
+# Now I finally implement A* and its most important helpers
+
+def expand_game_state(game_state: GameState) -> \
+ List[Tuple[int, int, int, int, int]]:
+ """Expand a game state tuple into a list of tuples specifying the room,
+ spot, home_room, amphipod index, and location of each amphipod,
+ reverse-sorted by room/spot"""
+ return sorted([(*location_to_room_spot(l), amphipod_type(a), a, l) \
+ for a, l in enumerate(game_state)],
+ reverse=True)
+
+def list_neighbors(game_state: GameState) -> List[Tuple[GameState, int]]:
+ """Return a list for almost every (legal) game state that can be reached from the
+ current one, as a tuple containing the state itself and the energy that
+ would be required to move there (there is some pruning here: if there are
+ paths that send amphipods home, only those neighbors are returned"""
+ # Before we start, we have to find all the _cozy_ amphipods and _ready_
+ # rooms. A cozy amphipod is in its home room and isn't blocking any
+ # amphipod that's not. A ready room is a room in which any amphipods that
+ # are present are cozy.
+ game_state_expanded = expand_game_state(game_state)
+ location_empty = [True] * (11 + ROOM_SIZE * 4)
+ amphipod_cozy = [False] * len(game_state)
+ room_ready = [False] + [True] * 4 # hallway is never ready
+ last_was_cozy = False
+ for room, spot, home_room, amphipod, location in game_state_expanded:
+ location_empty[location] = False
+ if room == home_room:
+ if spot == ROOM_SIZE - 1:
+ amphipod_cozy[amphipod] = True
+ last_was_cozy = True
+ #room_ready[room] = True
+ elif last_was_cozy:
+ amphipod_cozy[amphipod] = True
+ elif room > 0:
+ last_was_cozy = False
+ room_ready[room] = False
+
+ # Now we'll loop through all the amphipods again and figure out where each
+ # can move. We'll build:
+ # * A list of all neighbor nodes to the current game state
+ reachable_states = []
+ # * A subset of neighbors that put amphipods into rooms
+ reachable_room_states = []
+ # Once we've found a neighbor state that puts an amphipod into a room, we
+ # won't bother with any other type of route
+ route_to_room = False
+ for room, spot, home_room, amphipod, _ in game_state_expanded:
+ legal_moves = []
+ # We'll consider three situations in turn:
+ # * If the amphipod is in its home room, we'll move it to the hallway
+ # only if there is an amphipod of the wrong type below it (i.e., if
+ # it's not cozy).
+ # * If the amphipod is in the hallway, we'll move it to its home room
+ # only if any amphipods in the room are of the correct type (i.e, if
+ # the room is ready).
+ # * If the amphipod is in another room, we'll first check if there's a
+ # path to its home room (using the same rules as above); if not, then
+ # we'll consider moves into the hallway.
+ if room == home_room:
+ if not route_to_room and not amphipod_cozy[amphipod]:
+ # We need to consider moves to the hallway to get this amphipod
+ # out of the way of the amphipods below it (which need to move)
+ legal_moves = room_to_hallway(room, spot, location_empty)
+ elif room == 0:
+ if room_ready[home_room]:
+ # Consider moves to the home room (from the hallway)
+ legal_moves = hallway_to_room(spot, home_room, location_empty)
+ route_to_room = bool(legal_moves)
+ else:
+ if room_ready[home_room]:
+ legal_moves = room_to_room(room, spot, home_room,
+ location_empty)
+ route_to_room = bool(legal_moves)
+ if not route_to_room and not legal_moves:
+ # No route directly to home room, consider moves to the hallway
+ legal_moves = room_to_hallway(room, spot, location_empty)
+
+ for new_location, num_steps in legal_moves:
+ reachable_states.append((
+ make_state(game_state, amphipod, new_location),
+ get_energy_per_step(home_room) * num_steps
+ ))
+ if route_to_room:
+ reachable_room_states.append(reachable_states[-1])
+
+ if reachable_room_states:
+ return reachable_room_states
+ return reachable_states
+
+def heuristic_to_end(game_state: GameState) -> int:
+ """A heuristic that estimates the total energy used between a given game
+ state and an ending condition"""
+ # We'll loop through amphipods twice: first to find out which amphipods are
+ # cozy (and their locations, which we'll mark as unavailable).
+ game_state_expanded = expand_game_state(game_state)
+ cozy_amphipods = set()
+ available_locations = set(range(11, 11 + 4 * ROOM_SIZE))
+ last_was_cozy = False
+ for room, spot, home_room, amphipod, location in game_state_expanded:
+ if room == home_room:
+ if spot == ROOM_SIZE - 1:
+ cozy_amphipods.add(amphipod)
+ available_locations.remove(location)
+ last_was_cozy = True
+ elif last_was_cozy:
+ cozy_amphipods.add(amphipod)
+ available_locations.remove(location)
+ elif room > 0:
+ last_was_cozy = False
+
+ # Then, for the second loop, we'll put each amphipod into an available
+ # spot. Note that we're ignoring whether any other amphipods are in the way
+ # unless they're cozy.
+ energy = 0
+ for room, spot, home_room, amphipod, location in game_state_expanded:
+ if amphipod not in cozy_amphipods:
+ # Choose an available spot from their home room
+ for to_spot in range(ROOM_SIZE):
+ to_location = room_spot_to_location(home_room, to_spot)
+ if to_location in available_locations:
+ available_locations.remove(to_location)
+ break
+ else:
+ raise RuntimeError('no spot for amphipod')
+
+ # Find the steps/energy to this spot
+ num_steps = to_spot + 1
+ if room == 0:
+ hall_spot = spot
+ else:
+ hall_spot = room_entrance_locations()[room]
+ assert hall_spot is not None
+ num_steps += spot + 1
+ if room == home_room:
+ # This is a non-cozy amphipod in its own room, so it'll need to
+ # leave the room and return
+ num_steps += 2
+ num_steps += abs(hall_spot - room_entrance_locations()[home_room])
+ energy += num_steps * get_energy_per_step(home_room)
+ return energy
+
+def reconstruct_path(came_from: Dict[GameState, GameState],
+ current_state: GameState) -> Deque[GameState]:
+ """Return the step-by-step solution to the problem"""
+ # I didn't need to implement this to solve the puzzle, but I did to debug
+ # my solution. :(
+ # Also, a deque is overkill for this but I didn't use any for AoC and I
+ # wanted to (efficiently) appendleft so bad
+ total_path = deque([current_state])
+ while current_state in came_from:
+ current_state = came_from[current_state]
+ total_path.appendleft(current_state)
+ return total_path
+
+def find_lowest_energy_use(start_state: GameState) -> Tuple[GameState, int,
+ Deque[GameState]]:
+ """Use the A* algorithm to find the lowest energy use to get all the
+ amphipods from the starting game state to a completed game state"""
+ # Influenced heavily by the priority queue implementation in Python's heapq
+ # documentation. Note that passing multiple types into List[] is
+ # technically incorrect, but we're using lists instead of types
+ # specifically to get pass-by-reference behavior, so I'm accepting the mypy
+ # errors as a cost of clearer documentation.
+ open_set: List[List[int, Union[Tuple[int], GameState]]] = []
+ open_set_entries: Dict[GameState, List[int, GameState]] = {}
+
+ came_from: Dict[GameState, GameState] = {}
+ best_distance: Dict[GameState, Union[int, float]] = {start_state: 0}
+ best_cost: Dict[GameState, Union[int, float]] = {
+ start_state: 0 + heuristic_to_end(start_state)
+ }
+
+ current_state = start_state
+ while not everybody_home(current_state):
+ for neighbor, energy_between in list_neighbors(current_state):
+ energy_to_neighbor = best_distance[current_state] \
+ + energy_between
+ # If the distance to the neighbor from the start through the
+ # current node is lower than what was previously estimated, or if
+ # we've never seen a path to this neighbor before, update the
+ # priority queue and other data sctructures.
+ if energy_to_neighbor < best_distance.get(neighbor, inf):
+ estimated_cost = energy_to_neighbor \
+ + heuristic_to_end(neighbor)
+ came_from[neighbor] = current_state
+ best_distance[neighbor] = energy_to_neighbor
+ best_cost[neighbor] = estimated_cost
+ # Updating the priority queue is a bit more complicated
+ new_pq_entry = [estimated_cost, neighbor]
+ if neighbor in open_set_entries:
+ old_pq_entry = open_set_entries.pop(neighbor)
+ # Placeholder for "empty" priority queue entries
+ old_pq_entry[-1] = (0,)
+ open_set_entries[neighbor] = new_pq_entry
+ heappush(open_set, new_pq_entry)
+
+ # Select the next lowest-energy state to visit next.
+ while open_set:
+ _, current_state = heappop(open_set)
+ if current_state != (0,):
+ del open_set_entries[current_state]
+ break
+ else:
+ raise KeyError('pop from empty priority queue')
+
+ return (current_state, int(best_distance[current_state]),
+ reconstruct_path(came_from, current_state))
+
+def solve_puzzle(input_string: str) -> int:
+ """Return the numeric solution to the puzzle"""
+ return find_lowest_energy_use(parse_input(input_string))[1]
+
+def main() -> None:
+ """Run when the file is called as a script"""
+ assert solve_puzzle(EXAMPLE_INPUT) == 44169
+ print("Lowest energy use:",
+ solve_puzzle(get_puzzle_input(23)))
+
+if __name__ == "__main__":
+ main()
