Advent of Code 2021 - Day 23

Day 23: Amphipod

Useful Links

‹›

• 2021 Main Page
• My solution code

Concepts and Packages Demonstrated

Breadth First Search (BFS)Dijkstraheapqdictionary comprehensionzipany

class attributes__hash____eq____lt__ (<)

Problem Intro

Yep, this one was hard too. Not as hard as yesterday though. It was much easier to see how to do the solution. But I ended up writing quite a lot of code.

Overview

We’re told we have four types of amphipods, labelled A, B, C and D. They live in a burrow made up of a hallway and side rooms. Each room has two slots, and we start with all room slots filled with amphipods, using a configuration given as our starting input.

The sample input looks like this:

#############
#...........#
###B#C#B#D###
  #A#D#C#A#
  #########

A . char denotes an empty space.

Part 1

What is the least energy required, in order to rearrange the amphipods such that all the A pods are in the first room, all of the B pods are in the second, and so on?

So ultimately, we need to end up in this configuration:

#############
#...........#
###A#B#C#D###
  #A#B#C#D#
  #########

There are rules to amphipod movement:

• They can only move to an unoccuped space.
• They cannot move past another amphipod.
• They are not allowed to stop in the hall if they stop immediately above (and thus block) a room. Otherwise, any empty space in the hall is fair game.
• They are only allowed to move into a room, if that room is their correct target room. E.g. a B pod can only enter the second (B) room. A C pod can only enter the third (C) room.
• Furthermore, an amphipod will not enter their target room, if it has any ‘wrong’ amphipods in it. E.g. if the second room contains a B, then a B can enter it. But if the second room contains a C, then nothing can enter this room.
• If an amphipod stops in the hallway, it cannot move again, unless it is moving directly to its target room.
• Lastly, amphipod movement requires different amounts of energy:
  • A = 1 per step
  • B = 10 per step
  • C = 100 per step
  • D = 1000 per step

Thus, it’s intuitively obvious that D pods need to move as little as possible!

For me, there were two obvious ways to go about this:

1. With a Dijkstra breadth first search (BFS).
   • I.e. where we track the current burrow state.
   • We have a method that returns all possible next states.
   • We then use Dijkstra BFS to trawl through states from start to goal, but favouring moves with the lowest cost. (That’s what Dijkstra is good for!)
2. With recursion with memoization.
   • We use recursion to determine all possible next states, depth first.
   • We stop when we find a state that is the goal.
   • We use memoization to ensure that if we recurse into a state we’ve seen before, we immediately pop out the lowest cost answer that we’ve previously evaluated.

Personally, I like the Dijkstra approach more. So that’s what I’m going to describe here.

Let’s start with reading in the data:

input_file = os.path.join(SCRIPT_DIR, INPUT_FILE)
with open(input_file, mode="rt") as f:
    data = f.read().splitlines()

# Part 1
# Process data with initial state of our four rooms
room_a: list[str] = []
room_b: list[str] = []
room_c: list[str] = []
room_d: list[str] = []
hall: list[str] = []

for line in data:   # we only care about chars in "ABCD."
    for x, char in enumerate(line):
        if char == ".":     # hall
            hall.append(char)
        if char in "ABCD":  # rooms
            if x == 3:
                room_a.append(char)
            if x == 5:
                room_b.append(char)
            if x == 7:
                room_c.append(char)
            if x == 9:
                room_d.append(char) 

I’m representing the four rooms, and the hall, with lists that can only contain str.

• The hall will be a list with 11 items.
  • Each item can be either . or an amphipod type, A, B, C, or D.
  • Initially, the hall is empty. We rely on this when parsing the data. I.e. any . will be considerd a hall character.
• A room will be a list with two items.
  • Each item can be either . or an amphipod type, A, B, C, or D.
  • Initially, none of the rooms are empty. We rely on this when parsing the data. I.e. any char that is in "ABCD" will be considered a room character.

Now I create a dictionary that maps each room to the contents of that room. Each room key is one of A, B, C, or D, and each room value is the list that stores the room’s members.

# E.g. {'A': ['A', 'B'], 'B': ['D', 'C'], 'C': ['B', 'D'], 'D': ['C', 'A']}
rooms = {k: room for k, room in zip(BurrowState.ROOM_KEYS, [room_a, room_b, room_c, room_d])}

This works by:

• Using zip to zip our room keys (A, B, C, or D) to our room lists.
• Using dictionary comprehension to convert each tuple from the zip into a key:value dictionary item.

Now we create an initial BurrowState object, passing in a tuple that contains the initial values of rooms and hall:

start = BurrowState((rooms, hall))

As with all BFS and Dijkstra implementations, we’re going to need to compare our BurrowState objects, i.e. to see if we’ve seen this state before. My BurrowState class looks like this:

class BurrowState():
    """ Store the state of a Burrow, i.e. the configuration of the hall and rooms. 
    Knows how to yield next possible states (e.g. for use in a BFS). """
    A = 'A'
    B = 'B'
    C = 'C'
    D = 'D'
    EMPTY = '.'
    ROOM_KEYS = [A, B, C, D]
    POD_COSTS = {A: 1, B: 10, C: 100, D: 1000}
    ROOM_IDX = {A: 2, B: 4, C: 6, D: 8}

    def __init__(self, state: tuple[dict[str, list[str]], list[str]], cost:int=0) -> None:
        """ Creates a new BurrowState.

        Args:
            state (tuple[dict[str, list[str]], list[str]]): (rooms, hall)
            energy (int, optional): Energy required to get to this state. Defaults to 0.
        """
        self._state = state # (rooms, hall)
        self._last_cost = cost # energy required to get to this state from last state
    
    @property
    def last_cost(self):
        """ The energy cost taken to generate this state from the previous state.
        (This is not a cumulative energy.) """
        return self._last_cost
    
    def is_goal(self) -> bool:
        """ Returns False if any amphipods in any room are not of the right type """
        rooms, _ = self._state
        for room_type, pods in rooms.items():
            if any((amphipod != room_type) for amphipod in pods):
                return False
        return True

    def _can_move_from(self, room_key: str, room: list) -> bool:
        """ If this item is an amphipod and if we can move it """
        for item in room:
            if item != room_key and item != BurrowState.EMPTY:
                # If there's a pod at this location, and it doesn't belong in this room
                return True
        return False

    def _can_move_to(self, room_key: str, room: list) -> bool:
        """ Check if destination room is correct type and can accept a pod """
        for item in room:
            if item != room_key and item != BurrowState.EMPTY:
                # If there's already a pod in this room, and it's the wrong type 
                return False
        return True

    def _get_top_item_idx(self, room_contents: list):
        """ Return the row index of the top pod (i.e. top occupied position) in a room """
        for i, item in enumerate(room_contents):
            if item != BurrowState.EMPTY:
                return i
        return None

    def _get_room_dest_idx(self, room_contents: list):
        """ Return the position in the room we want to move to. """
        for i, char in reversed(list(enumerate(room_contents))):
            if char == BurrowState.EMPTY:
                return i    # return the "bottom" (highest index) that is empty
        return None

    def _is_between(self, posn: int, room_key: str, hall_idx: int) -> bool:
        """ If this posn is between the room and the hall index """
        return ((BurrowState.ROOM_IDX[room_key] < posn < hall_idx)
                or (hall_idx < posn < BurrowState.ROOM_IDX[room_key]))

    def _is_clear_path(self, room_key: str, hall_idx: int) -> bool:
        """ Is it clear between the room and the hall position? 

        Args:
            room_key (str): Which room
            hall_idx (int): Which hall horizontal position
            hall (list[str]): Contents of the hall
        """
        _, hall = self._state
        
        for posn, item in enumerate(hall):
            if self._is_between(posn, room_key, hall_idx) and item != BurrowState.EMPTY:
                return False
        return True

    def __repr__(self) -> str:
        """ Generate a str representation of this state """
        rooms, hall = self._state
        rooms_list = [room_items for room_type, room_items in rooms.items()]
        render = []
        render.append('')  # Blank line
        render.append("#" + "#"*len(hall) + "#") # top row
        render.append("#" + "".join(hall) + "#") # hall row
        for i in range(len(rooms_list[0])): # room rows
            if i == 0:
                prefix = suffix = "###" # top room row
            else:
                prefix = "  #"
                suffix = "#"
                
            render.append(prefix + "#".join(rooms[k][i] for k in rooms) + suffix)
        render.append("  " + "#"*(len(hall)-2)) # bottom row
        return "\n".join(render)

    def __hash__(self) -> int:
        rooms, hall = self._state
        rooms_tuple = tuple((k, tuple(v)) for k,v in rooms.items())
        hall_tuple = tuple(hall)
        return hash(tuple((rooms_tuple, hall_tuple)))
    
    def __eq__(self, __o: object) -> bool:
        """ Test for equivalence by comparing hall and rooms """
        if isinstance(__o, BurrowState):
            this_rooms, this_hall = self._state
            other_rooms, other_hall = __o._state
            
            return this_rooms == other_rooms and this_hall == other_hall
        else:
            return NotImplemented
        
    def __lt__(self, __o: object) -> bool:
        """ Performs comparision based on last energy cost. """
        if isinstance(__o, BurrowState):
            return self.last_cost < __o.last_cost
        else:
            return NotImplemented
    
    def next_state(self) -> Iterable[BurrowState]:
        """ Yields next possible states from here.
        If we can put a amphipod into its destination room, then only return that state. """  
        rooms, hall = self._state
    
        # Determine if we have any pods in the hall that can go directly to dest
        # If we have, then this is the best move
        for i, pod in enumerate(hall): # position and item (pod or empty) in hall
            # If a pod, and we can move it to the room
            if pod in rooms and self._can_move_to(pod, rooms[pod]):
                if self._is_clear_path(pod, i):
                    dest_idx = self._get_room_dest_idx(rooms[pod])
                    assert isinstance(dest_idx, int), "We've determined we can move to this room"
                    dist = dest_idx + 1 + abs(BurrowState.ROOM_IDX[pod]-i)
                    cost = BurrowState.POD_COSTS[pod] * dist
                    
                    # remove pod from hall
                    new_hall = list(hall) # create a copy
                    new_hall[i] = BurrowState.EMPTY
                    
                    # add pod to destination room
                    new_rooms = {k: [room_item for room_item in v] for k, v in rooms.items()}                
                    new_rooms[pod][dest_idx] = pod

                    yield BurrowState((new_rooms, new_hall), cost=cost)

        # If we're here, no pods in the hall to move to destination.
        # Evaluate which rooms we can move pods from
        for room_key, room_contents in rooms.items():
            if not self._can_move_from(room_key, room_contents):
                continue
            # If we're here, we can move out of this room
            
            pod_idx = self._get_top_item_idx(room_contents)
            if pod_idx is None:
                continue
            
            # If we're here, there's a pod to move from this room
            pod = room_contents[pod_idx]
            
            # Now find locations in the hall our pod can move to
            for hall_posn, item in enumerate(hall):
                if hall_posn in [2, 4, 6, 8]:
                    continue    # skip hall positions above rooms
                if item != BurrowState.EMPTY:
                    continue    # skip locations that are occupied
                
                # determine if we have a path to this destination
                if self._is_clear_path(room_key, hall_posn):
                    dist = pod_idx + 1 + abs(hall_posn - BurrowState.ROOM_IDX[room_key])
                    cost = BurrowState.POD_COSTS[pod] * dist  # cost of this move
                    
                    # make a copy of hall and rooms, and update them
                    new_hall = list(hall)
                    new_hall[hall_posn] = pod
                    
                    # This is much quicker than a deep copy
                    new_rooms = {k: [room_item for room_item in v] for k, v in rooms.items()}
                    new_rooms[room_key][pod_idx] = BurrowState.EMPTY
                    
                    yield BurrowState((new_rooms, new_hall), cost=cost)

Yeah, it is quite long. Sorry! We could be briefer, but I’ve gone for readability.

Let’s break it down.

First, some class attributes, that we can treat as constants:

    A = 'A'
    B = 'B'
    C = 'C'
    D = 'D'
    EMPTY = '.'
    ROOM_KEYS = [A, B, C, D]
    POD_COSTS = {A: 1, B: 10, C: 100, D: 1000}
    ROOM_IDX = {A: 2, B: 4, C: 6, D: 8}

Note that ROOM_IDX is a dictionary, where the key is the room, and the value is the index of where that room is found, relative to hall. This is clearer if I replace some hall spaces with index values in the layout diagram:

#############
#012345678..#
###A#B#C#D###
  #A#B#C#D#
  #########

Next, the initialiser:

    def __init__(self, state: tuple[dict[str, list[str]], list[str]], cost:int=0) -> None:
        """ Creates a new BurrowState.

        Args:
            state (tuple[dict[str, list[str]], list[str]]): (rooms, hall)
            energy (int, optional): Energy required to get to this state. Defaults to 0.
        """
        self._state = state # (rooms, hall)
        self._last_cost = cost # energy required to get to this state from last state

As we’ve already seen, we can create a BurrowState object by passing in a tuple of (rooms, hall). The tuple is stored in the instance’s _state attribute.

We also store an attribute called _last_cost. This stores the cost that was required for the last move. And consequently, our initialiser has a third optional argument, which is the cost of the last move made that resulted in this BurrowState. This defaults to 0, reflecting the fact that the very first BurrowState we create will have no cost. But all subsequent BurrowState objects will have a cost, i.e. because an amphipod will have moved.

Now let’s look at some of the methods:

    def is_goal(self) -> bool:
        """ Returns False if any amphipods in any room are not of the right type """
        rooms, _ = self._state
        for room_type, pods in rooms.items():
            if any((amphipod != room_type) for amphipod in pods):
                return False
        return True

    def _can_move_from(self, room_key: str, room: list) -> bool:
        """ If this item is an amphipod and if we can move it """
        for item in room:
            if item != room_key and item != BurrowState.EMPTY:
                # If there's a pod at this location, and it doesn't belong in this room
                return True
        return False

    def _can_move_to(self, room_key: str, room: list) -> bool:
        """ Check if destination room is correct type and can accept a pod """
        for item in room:
            if item != room_key and item != BurrowState.EMPTY:
                # If there's already a pod in this room, and it's the wrong type 
                return False
        return True

    def _get_top_item_idx(self, room_contents: list):
        """ Return the row index of the top pod (i.e. top occupied position) in a room """
        for i, item in enumerate(room_contents):
            if item != BurrowState.EMPTY:
                return i
        return None

    def _get_room_dest_idx(self, room_contents: list):
        """ Return the position in the room we want to move to. """
        for i, char in reversed(list(enumerate(room_contents))):
            if char == BurrowState.EMPTY:
                return i    # return the "bottom" (highest index) that is empty
        return None

    def _is_between(self, posn: int, room_key: str, hall_idx: int) -> bool:
        """ If this posn is between the room and the hall index """
        return ((self._get_room_horiz_idx(room_key) < posn < hall_idx)
                or (hall_idx < posn < self._get_room_horiz_idx(room_key)))

    def _is_clear_path(self, room_key: str, hall_idx: int) -> bool:
        """ Is it clear between the room and the hall position? 

        Args:
            room_key (str): Which room
            hall_idx (int): Which hall horizontal position
            hall (list[str]): Contents of the hall
        """
        _, hall = self._state
        
        for posn, item in enumerate(hall):
            if self._is_between(posn, room_key, hall_idx) and item != BurrowState.EMPTY:
                return False
        return True

• last_cost is a convenience property, to expose the cost to any calling code.
• is_goal() determines if this BurrowState is the desired final state. It does this by iterating through items of each room and checking if any item in the room is not the required type of amphipod. (Empty spaces would count as “not the right type of amphipod”.) Note the use of the any() function here.
  • This evaluates a list of booleans, and returns True if any members of the list evaluate to True.
  • In this case, we use list comprehension to generate a list of boolean values, based on checking whether any item in the room is of the wrong type. If any item is of the wrong type, then any() will evaluate to True. Thus, if any items in the room are wrong type, then we haven’t reached the goal, so our method returns False.
• _can_move_from() is a private method that takes a room and its items, and returns True if the room contains an amphipod that is wrong type. (Otherwise, the amphipod should stay put!)
• _can_move_to() is a private method that takes a room and its items, and checks whether the room contains any amphipods of the wrong type. If it does, it cannot accept any amphipods, as per the rules!
• _get_top_item_idx() is a private method that takes the items in any room, and retuns the vertical index value of the first amphipod found in that room, counting from the top. Thus, if the room contains two amiphods, then the top index is 0. If it contains one amphipod, then the top index is 1. If it contains no amphipods, then the method returns None.
• _get_room_dest_idx() is a private method that takes the items of a given room, and returns the vertical index value of the first emtpy space found in the room, counting from the bottom. I.e. this is finding the index of the first position we can move an amphipod to.
• _is_between() is a private method that determines if the specified position is between a given room position, and a specified hall position.
• _is_clear_path() is a private method that relies on _is_between() to determine if there are any amphipods between a room position, and a specified hall position. It works by working through each position along the hall, and if that position _is_between() the room position and the hall position, and that position is not empty, then the path is blocked.

Now we need to be able to compare BurrowState objects, to check if we’ve seen this BurrowState before. As always, in order to check for equality, we need to implement the __hash__() method, as well as the __eq__() method. In addition, in order to be able to push and pop BurrowState objects on/off a priority queue, we must also implement the __lt__(), which provides the implementation for the < operator.

    def __hash__(self) -> int:
        rooms, hall = self._state
        rooms_tuple = tuple((k, tuple(v)) for k,v in rooms.items())
        hall_tuple = tuple(hall)
        return hash(tuple((rooms_tuple, hall_tuple)))
    
    def __eq__(self, __o: object) -> bool:
        """ Test for equivalence by comparing hall and rooms """
        if isinstance(__o, BurrowState):
            this_rooms, this_hall = self._state
            other_rooms, other_hall = __o._state
            
            return this_rooms == other_rooms and this_hall == other_hall
        else:
            return NotImplemented
        
    def __lt__(self, __o: object) -> bool:
        """ Performs comparision based on last energy cost. """
        if isinstance(__o, BurrowState):
            return self.last_cost < __o.last_cost
        else:
            return NotImplemented

• The __hash__() method works by creating tuples from the room contents and hall contents, then hashing the tuple of these two tuples. We do this because the hash() method only works on immutable objects, and converting something to a tuple is a convenient way to make something immutable.
• The __eq__() method compares this object to another object. If the other object is not a BurrowState, then return NotImplemented. Otherwise, compare the rooms and the hall, and see if they are the same.
• The __lt__() method is similar, but compares the last_cost property. This is because when we do our Dijktra BFS, we want to compare costs of getting to a given state. We always want to prioritise lowest cost.

Now let’s implement a method to generate a nice visual representation of any BurrowState object:

    def __repr__(self) -> str:
        """ Generate a str representation of this state """
        rooms, hall = self._state
        rooms_list = [room for room_type, room in rooms.items()]
        render = []
        render.append('')  # Blank line
        render.append("#" + "#"*len(hall) + "#") # top row
        render.append("#" + "".join(hall) + "#") # hall row
        for i in range(len(rooms_list[0])): # room rows
            if i == 0:
                prefix = suffix = "###" # top room row
            else:
                prefix = "  #"
                suffix = "#"
                
            render.append(prefix + "#".join(rooms[k][i] for k in rooms) + suffix)
        render.append("  " + "#"*(len(hall)-2)) # bottom row
        return "\n".join(render)

This works as follows:

• Get a list of all room items.
• Create a render list, that we’ll build up as we go.
  • Add in the top row of # by:
    • Adding a prefix #.
    • Adding "#"*len(hall) - i.e. the number of hashes equivelent to the length of the hall. It’s cool how we can multiply strings, right?
    • Adding a final #.
  • Add in the hall row by:
    • Adding a prefix #.
    • Using "".join(hall) to convert the hall list into a str.
    • Adding a final #.
  • Now add in the room rows, by turning the room items to a str, and joining them with #, with prefix and suffix # characters as required.
  • Add in the bottom row of #.
• Finally, convert our render list into a multiline str, by using "\n".join(render).

And the final method of our BurrowState is the next_state() method, which returns all possible next BurrowState objects that can be reached in one move:

    def next_state(self) -> Iterable[BurrowState]:
        """ Yields next possible states from here.
        If we can put a amphipod into its destination room, then only return that state. """  
        rooms, hall = self._state
    
        # Determine if we have any pods in the hall that can go directly to dest
        # If we have, then this is the best move
        for i, pod in enumerate(hall): # position and item (pod or empty) in hall
            # If a pod, and we can move it to the room
            if pod in rooms and self._can_move_to(pod, rooms[pod]):
                if self._is_clear_path(pod, i):
                    dest_idx = self._get_room_dest_idx(rooms[pod])
                    assert isinstance(dest_idx, int), "We've determined we can move to this room"
                    dist = dest_idx + 1 + abs(BurrowState.ROOM_IDX[pod]-i)
                    cost = BurrowState.POD_COSTS[pod] * dist
                    
                    # remove pod from hall
                    new_hall = list(hall) # create a copy
                    new_hall[i] = BurrowState.EMPTY
                    
                    # add pod to destination room
                    new_rooms = {k: [room_item for room_item in v] for k, v in rooms.items()}                
                    new_rooms[pod][dest_idx] = pod

                    yield BurrowState((new_rooms, new_hall), cost=cost)

        # If we're here, no pods in the hall to move to destination.
        # Evaluate which rooms we can move pods from
        for room_key, room_contents in rooms.items():
            if not self._can_move_from(room_key, room_contents):
                continue
            # If we're here, we can move out of this room
            
            pod_idx = self._get_top_item_idx(room_contents)
            if pod_idx is None:
                continue
            
            # If we're here, there's a pod to move from this room
            pod = room_contents[pod_idx]
            
            # Now find locations in the hall our pod can move to
            for hall_posn, item in enumerate(hall):
                if hall_posn in [2, 4, 6, 8]:
                    continue    # skip hall positions above rooms
                if item != BurrowState.EMPTY:
                    continue    # skip locations that are occupied
                
                # determine if we have a path to this destination
                if self._is_clear_path(room_key, hall_posn):
                    dist = pod_idx + 1 + abs(hall_posn - BurrowState.ROOM_IDX[room_key])
                    cost = BurrowState.POD_COSTS[pod] * dist  # cost of this move
                    
                    # make a copy of hall and rooms, and update them
                    new_hall = list(hall)
                    new_hall[hall_posn] = pod
                    
                    # This is much quicker than a deep copy
                    new_rooms = {k: [room_item for room_item in v] for k, v in rooms.items()}
                    new_rooms[room_key][pod_idx] = BurrowState.EMPTY
                    
                    yield BurrowState((new_rooms, new_hall), cost=cost)

This is where most of the heavy lifting happens.

• First, note that there are only two possible moves we consider:
  1. An amphipod moving from the hall to a room.
  2. An amphipod moving from a room to the hall.
• We start by looking for any pods in the hall that can go directly to their target room. If so, then this is the best next move.
  • Iterate through each position in the hall.
  • If the item at that position is an ampiphod (which we can check by seeing if the item matches a room key), then:
    • If its respective room can accept it AND there’s a clear path to the room, then:
      • Update the hall configuration to reflect the move. (I.e. this position will now be empty.)
      • Update the room configuration to reflect the move. (I.e. this spot in the room will now contain the amphipod.)
      • Determine the distance the amphipod moved.
      • Determine the cost of the move.
      • Generate a new BurrowState from the current (rooms, hall), and pass in the cost for this amphipod move. Yield this BurrowState.
• If there were no pods in the hall that could move to a room, then we need to move a pod out of a room.
  • Iterate through all rooms.
    • For each, determine if a pod can be moved out of this room.
    • Get the ‘top’ pod in the room.
    • Determine which hall positions the pod can move to. It must be an empty position, and not blocking a room.
    • Determine if there’s a clear path to that position in the hall. If so:
      • Update the hall configuration to reflect the move. (I.e. this position will now contain the amphipod.)
      • Update the room configuration to reflect the move. (I.e. this spot in the room will now be empty.)
      • Determine the distance the amphipod moved.
      • Determine the cost of the move.
      • Generate a new BurrowState from the current (rooms, hall), and pass in the cost for this amphipod move. Yield this BurrowState.

Finally, we’re ready to implement our Dijkstra function:

def solve_with_dijkstra(start: BurrowState) -> tuple[BurrowState, dict[BurrowState, BurrowState]]:
    current: BurrowState = start
    frontier: list = []
    heapq.heappush(frontier, (0, current))   # init state will have energy cost of 0
    
    came_from = {}  # so we can rebuild path from breadcrumbs
    came_from[current] = None
    
    energy_so_far = {}  # store cumulative energy required to get to this state. Use as priority for heapq.
    energy_so_far[current] = 0
    
    while frontier:
        _, current = heapq.heappop(frontier)
        if current.is_goal():
            break
        
        next_state: BurrowState
        for next_state in current.next_state():
            new_energy = energy_so_far[current] + next_state.last_cost
            
            # If we haven't seen this state before, or we've found a more efficient way to get to this state...
            if next_state not in energy_so_far or new_energy < energy_so_far[next_state]:
                energy_so_far[next_state] = new_energy
                heapq.heappush(frontier, (new_energy, next_state))
                came_from[next_state] = current
                
    return (current, came_from)

This is what it does:

• Create our frontier, using a heapq. Recall that when we pop from the heapq, the lowest priority item is popped first.
• Add our initial state to the frontier, in the form of a tuple of (cost, BurrowState). Note that we have to put cost first, since when we’re popping, the heapq compares all the tuples it contains, and compares the first item in these tuples to determine lowest priority. When adding our initial state, the cost is 0. There can be no lower cost than this.
• We create a came_from dictionary, which we’ll use to point a given BurrowState object to its previous BurrowState. This will allow us to later create a “breadcrumb trail” from the final BurrowState, back to the first.
• We create an energy_so_far dictionary, which stores the cumulative energy required to reach a given BurrowState. Thus, if we see the same state more than once, we’ll keep it if its cost was lower than the previous route. And if it’s not, we just bin in and move on.
• Now we enter the familiar territory of popping off the frontier.
  • If the BurrowState popped is our goal, then we break.
  • Otherwise, we use next_state() on the current BurrowState, to iterate through the possible next states. For each next state:
    • We determine the cumulative cost, by adding the current cost to the cost of the last move.
    • If we’ve not seen this next state before, or we have seen this next state but this route has a lower cost, then we store this new state in the energy_so_far dictionary, add it to the came_from dict, and add it to the frontier.
• The solve_with_dijkstra() method returns a tuple made up of:
  • The final BurrowState object, which will be the goal.
  • The dictionary that maps all BurrowState objects to their previous states.

We call it like this:

current, came_from = solve_with_dijkstra(start)

So at this point, we have a dictionary of BurrowState objects, with each key pointing to the prevoius BurrowState. And because we’ve solved using Dijkstra, the sequence of keys will the lowest cost way to solve the problem.

But we still haven’t determined the cumulative cost. And that’s what we need. So to finally get to the answer, I’ve implemented a render_breadcrumbs() method.

def render_breadcrumbs(last_state: BurrowState, came_from: dict, start: BurrowState) -> str:
    """ Given a final state and a map of state-to-previous state, 
    render the entire path to the initial state. """
    breadcrumbs = []    # BurrowStates
    current: BurrowState = last_state  
    cumulative_energy = 0
    while current != start:
        breadcrumbs.append(str(current))
        cumulative_energy += current.last_cost
        current = came_from[current]
    
    breadcrumbs.append(str(start))
    breadcrumbs.reverse()
    breadcrumbs.append(f"\nCompleted in {len(breadcrumbs)-1} steps with total energy of {cumulative_energy}.\n")
    return "\n".join(breadcrumbs)

This works by rebuilding the all the BurrowState objects, from final state back to the first. We iterate through each BurrowState, starting with the last, and add its last_cost to a cumulative_energy variable. Also, with each iteration, we add the str representation of the current BurrowState to a breadcrumbs list.

When we’ve finished iterating, we call reverse() on breadcrumbs, so that our list now goes from first to last BurrowState. Finally, we print the cumulative score, and we join all the breadcrumb strings together using "\n".join(breadcrumbs).

We call it like this:

logger.info("Part 1:\n%s", render_breadcrumbs(current, came_from, start))

It works. Hurrah!

Part 2

Now the rooms are bigger! We’re told that we need to add two additional starting lines between the existing room lines:

  #D#C#B#A#
  #D#B#A#C#

Once again, determine the least energy required, in order to rearrange the amphipods such that all the A pods are in the first room, all of the B pods are in the second, and so on.

Good news. I don’t need to change the solution at all!

It’ll just work, because we haven’t hard-coded anything about the size of the rooms. So all we have to do is modify our starting configuration - by adding in the two new rows - and then executing exactly as we did for Part 1.

So, we execute Part 2 like this:

# Part 2 - We need to insert D C B A
#                            D B A C
room_a.insert(1, 'D')
room_a.insert(2, 'D')
room_b.insert(1, 'C')
room_b.insert(2, 'B')
room_c.insert(1, 'B')
room_c.insert(2, 'A')
room_d.insert(1, 'A')
room_d.insert(2, 'C')

start = BurrowState((rooms, hall))
current, came_from = solve_with_dijkstra(start)
logger.info("Part 2:\n%s", render_breadcrumbs(current, came_from, start))

And that’s it.

It works. WOOP! Takes about 7 seconds to run.

Here’s the output using my data:

19:23:03.339:INFO:__main__:     Part 1:

#############
#...........#
###A#D#B#C###
  #B#C#D#A#
  #########

#############
#.A.........#
###.#D#B#C###
  #B#C#D#A#
  #########

#############
#.A.B.......#
###.#D#.#C###
  #B#C#D#A#
  #########

#############
#.A.B.C.....#
###.#D#.#.###
  #B#C#D#A#
  #########

#############
#.A.B.C...A.#
###.#D#.#.###
  #B#C#D#.#
  #########

#############
#.A.B.C.D.A.#
###.#D#.#.###
  #B#C#.#.#
  #########

#############
#.A.B...D.A.#
###.#D#.#.###
  #B#C#C#.#
  #########

#############
#.A.B.....A.#
###.#D#.#.###
  #B#C#C#D#
  #########

#############
#.A.B...D.A.#
###.#.#.#.###
  #B#C#C#D#
  #########

#############
#.A.B.C.D.A.#
###.#.#.#.###
  #B#.#C#D#
  #########

#############
#.A...C.D.A.#
###.#.#.#.###
  #B#B#C#D#
  #########

#############
#.A.B.C.D.A.#
###.#.#.#.###
  #.#B#C#D#
  #########

#############
#...B.C.D.A.#
###.#.#.#.###
  #A#B#C#D#
  #########

#############
#.....C.D.A.#
###.#B#.#.###
  #A#B#C#D#
  #########

#############
#.......D.A.#
###.#B#C#.###
  #A#B#C#D#
  #########

#############
#.........A.#
###.#B#C#D###
  #A#B#C#D#
  #########

#############
#...........#
###A#B#C#D###
  #A#B#C#D#
  #########

Completed in 16 steps with total energy of 13336.

19:23:07.569:INFO:__main__:     Part 2:

#############
#...........#
###A#D#B#C###
  #D#C#B#A#
  #D#B#A#C#
  #B#C#D#A#
  #########

#############
#..........B#
###A#D#.#C###
  #D#C#B#A#
  #D#B#A#C#
  #B#C#D#A#
  #########

#############
#.........BB#
###A#D#.#C###
  #D#C#.#A#
  #D#B#A#C#
  #B#C#D#A#
  #########

#############
#A........BB#
###A#D#.#C###
  #D#C#.#A#
  #D#B#.#C#
  #B#C#D#A#
  #########

#############
#AD.......BB#
###A#.#.#C###
  #D#C#.#A#
  #D#B#.#C#
  #B#C#D#A#
  #########

#############
#AD.D.....BB#
###A#.#.#C###
  #D#C#.#A#
  #D#B#.#C#
  #B#C#.#A#
  #########

#############
#AD.D.C...BB#
###A#.#.#C###
  #D#.#.#A#
  #D#B#.#C#
  #B#C#.#A#
  #########

#############
#AD.D.....BB#
###A#.#.#C###
  #D#.#.#A#
  #D#B#.#C#
  #B#C#C#A#
  #########

#############
#AD.D...B.BB#
###A#.#.#C###
  #D#.#.#A#
  #D#.#.#C#
  #B#C#C#A#
  #########

#############
#AD.D.C.B.BB#
###A#.#.#C###
  #D#.#.#A#
  #D#.#.#C#
  #B#.#C#A#
  #########

#############
#AD.D...B.BB#
###A#.#.#C###
  #D#.#.#A#
  #D#.#C#C#
  #B#.#C#A#
  #########

#############
#AD.D.....BB#
###A#.#.#C###
  #D#.#.#A#
  #D#.#C#C#
  #B#B#C#A#
  #########

#############
#AD.D......B#
###A#.#.#C###
  #D#.#.#A#
  #D#B#C#C#
  #B#B#C#A#
  #########

#############
#AD.D.......#
###A#.#.#C###
  #D#B#.#A#
  #D#B#C#C#
  #B#B#C#A#
  #########

#############
#AD.D...C...#
###A#.#.#.###
  #D#B#.#A#
  #D#B#C#C#
  #B#B#C#A#
  #########

#############
#AD.D...C..A#
###A#.#.#.###
  #D#B#.#.#
  #D#B#C#C#
  #B#B#C#A#
  #########

#############
#AD.D......A#
###A#.#.#.###
  #D#B#C#.#
  #D#B#C#C#
  #B#B#C#A#
  #########

#############
#AD.D...C..A#
###A#.#.#.###
  #D#B#C#.#
  #D#B#C#.#
  #B#B#C#A#
  #########

#############
#AD.D...C.AA#
###A#.#.#.###
  #D#B#C#.#
  #D#B#C#.#
  #B#B#C#.#
  #########

#############
#AD.D.....AA#
###A#.#C#.###
  #D#B#C#.#
  #D#B#C#.#
  #B#B#C#.#
  #########

#############
#AD.......AA#
###A#.#C#.###
  #D#B#C#.#
  #D#B#C#.#
  #B#B#C#D#
  #########

#############
#A........AA#
###A#.#C#.###
  #D#B#C#.#
  #D#B#C#D#
  #B#B#C#D#
  #########

#############
#AA.......AA#
###.#.#C#.###
  #D#B#C#.#
  #D#B#C#D#
  #B#B#C#D#
  #########

#############
#AA.....D.AA#
###.#.#C#.###
  #.#B#C#.#
  #D#B#C#D#
  #B#B#C#D#
  #########

#############
#AA...D.D.AA#
###.#.#C#.###
  #.#B#C#.#
  #.#B#C#D#
  #B#B#C#D#
  #########

#############
#AA.B.D.D.AA#
###.#.#C#.###
  #.#B#C#.#
  #.#B#C#D#
  #.#B#C#D#
  #########

#############
#A..B.D.D.AA#
###.#.#C#.###
  #.#B#C#.#
  #.#B#C#D#
  #A#B#C#D#
  #########

#############
#...B.D.D.AA#
###.#.#C#.###
  #.#B#C#.#
  #A#B#C#D#
  #A#B#C#D#
  #########

#############
#.....D.D.AA#
###.#B#C#.###
  #.#B#C#.#
  #A#B#C#D#
  #A#B#C#D#
  #########

#############
#.....D...AA#
###.#B#C#.###
  #.#B#C#D#
  #A#B#C#D#
  #A#B#C#D#
  #########

#############
#.........AA#
###.#B#C#D###
  #.#B#C#D#
  #A#B#C#D#
  #A#B#C#D#
  #########

#############
#..........A#
###.#B#C#D###
  #A#B#C#D#
  #A#B#C#D#
  #A#B#C#D#
  #########

#############
#...........#
###A#B#C#D###
  #A#B#C#D#
  #A#B#C#D#
  #A#B#C#D#
  #########

Completed in 32 steps with total energy of 53308.

19:23:07.642:INFO:__main__:     Execution time: 7.9547 seconds

‹›