My attempts to work through the 2021 Advent of Code problems.

```commit 50189f6b23b0e0553d41079ba54cc3f588cad4d6
parent 2f41fa16e2b71497ab2599939664efe08cff03ee
Date:   Mon, 27 Dec 2021 12:47:14 -0500

Solution to day 23, part 2

Diffstat:
```
```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
+            elif last_was_cozy:
+                amphipod_cozy[amphipod] = True
+        elif room > 0:
+            last_was_cozy = 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:
+                # 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:
+                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:
+                available_locations.remove(location)
+                last_was_cozy = True
+            elif last_was_cozy:
+                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)
+
+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()
```