Amphipod

Published: 2022-01-22 Original Prompt

Part 1

This has been the toughest day for me so far. I usually do both parts in a couple of hours, max. I worked on this one for literal days. I got there! But boy, did it take some time.

Generically, it’s the same problem as day 15- we have to figure out the most efficient bath between a starting state and and an ending state. We will be able to calculate the cost between two states. So, we’ll be able to use Dijkstra’s algorithm again to find the most efficient series of intermediate states to get to the desired ending.

To that end, I copied my entire algorithm from day 15 and put it in a self._dijkstra method. Its only parameter is max_room_size, which will be important for part 2; don’t worry about it for now.

The part of this that gave me the most trouble was how to model the data. It’s sort of a grid, but you can’t move freely. There are “houses” (vertical parts) and the “hallway” (horizontal area). An amphipod is “home” if it’s in the correct house. I first tried to organize it with nested lists, like so:

# the state in the example
[
    None,
    None,
    [None, "B", "A"],
    None,
    [None, "C", "D"],
    None,
    ...
]

Each step is either empty or has a letter; each room is a sub array. But, I found this a little hard to visualize and manipulate, so I went for a more traditional grid. Each point is a tuple of (horiz, vert). Here’s the initial example again:

(0,0) (1,0) (2,0) (3,0) (4,0) (5,0) ... (9,0) (10,0)
            (2,1)       (4,1)
            (2,2)       (4,2)

It’s a “grid”, but there’s a fixed number of valid points. For instance, (1,1) is a wall and can’t have an amphipod there, so we ignore it. This lets us store our state in a Dict[Point, str]. Let’s make our basic state class and helper types:

from dataclasses import dataclass
# we're not using these all yet, but we will!
from typing import (
    Callable,
    DefaultDict,
    Dict,
    List,
    Literal,
    Optional,
    Set,
    Tuple,
    Union,
    cast,
)

Amphipod = Union[Literal["A"], Literal["B"], Literal["C"], Literal["D"]]
GridPoint = Tuple[int, int]

@dataclass
class State:
    populated: Dict[GridPoint, Amphipod]
    max_room_size: int = 2

Now that we have somewhere to put them, we’ll need to parse our input:

class Solution:
    ...
    def _parse_input(
    start_order: List[Amphipod] = re.findall(r"[ABCD]", self.input)
        locations = {}
        i = 0
        for vert in [1, 2]:
            for horiz in [2, 4, 6, 8]:
                locations[(horiz, vert)] = start_order[i]
                i += 1
        return locations

That works because the regex finds letters left-to-right, top-to-bottom. We can assign them locations the same way.

Ok! So put those together and we have a state which knows exactly where all of the 8 amphipods are. The first thing we’ll need to be able to calculate is whether or not that state has “won” (so we know when to exit). A state has won if all of its amphipods are in their home columns. So, we’ll need a function that tells us that, too. They’re not too complicated:

# which column each amph belongs to
HOME_HORIZ: Dict[Amphipod, int] = {"A": 2, "B": 4, "C": 6, "D": 8}

class State:
    ...

    def loc_is_home(self, loc: GridPoint) -> bool:
        """
        Is the amphipod at the given location home?
        """
        if loc not in self.populated:
            return False

        return loc[1] > 0 and loc[0] == HOME_HORIZ[self.populated[loc]]

    @property
    def did_win(self) -> bool:
        return all(self.loc_is_home(p) for p in self.populated.keys())

Now, given a state, we know if it’s won. We’ve confirmed our initial state hasn’t won (bummer, that would save a lot of time here). Now we need to be able to get a new state from an old state. The new state is mostly identical to the old one, but its populated dict will have exactly one key different:

class State:
    ...

    def new_state_with_swap(self, old: GridPoint, new: GridPoint) -> "State":
        new_pop = self.populated.copy()
        del new_pop[old]
        new_pop[new] = self.populated[old]
        return State(new_pop, self.max_room_size)

Turn it around and that’s a brand new State! In practice, we’ll be coupling a new state with how much it cost to transition between them (that cost is very important to our algorithm). Given those two points, we can calculate their distance as the horizontal distance plus each of their vertical distances (which covers cases where we have to go up, side, and then down):

COST: Dict[Amphipod, int] = {"A": 1, "B": 10, "C": 100, "D": 1000}

class State:
    ...

    def cost_between_points(self, a: GridPoint, b: GridPoint) -> int:
        # exactly one of these points will have an amphipod; you can verify that by adding
        # an `assert self.populated.get(a) ^ self.populated.get(b)`
        # which uses Python's xor operator to verify our requirement
        step_cost = COST[cast(Amphipod, self.populated.get(a) or self.populated.get(b))]
        return (a[1] + b[1] + abs(a[0] - b[0])) * step_cost

There’s a bit of extra typing cruft (the cast) in there, but the math is on the last line and is straightforward. We can wrap these two methods into a single function that returns the 2-tuple we’ll eventually be sorting in Dijkstra’s:

class State:
    ...

    def _new_state(self, loc: GridPoint, to: GridPoint) -> Tuple[int, "State"]:
        return (
            self.cost_between_points(loc, to),
            self.new_state_with_swap(loc, to),
        )

This all looks good, but there’s a catch- nothing we’ve written so far verifies the validity of the moves. Our prompt explains those rules, and there are many cases where we’re unable to move an amphipod between two points (e.g. it’s path is blocked). Let’s think through the rules of how an amphipod can move. Then, given a state, we’ll be able to generate every possible next state. There are a few rules of motion to consider:

If an amphipod is in a house it can:
1. move into the hallway
2. move home
If an amphipod is in the hallway, it can only move home
1. To move home, all amphipods in that house must already be home
2. To move home, there must be a clear path there
If an amphipod is in a house and is not at the top, it can’t move
An amphipod will only make 1 or 2 moves during the game:
1. start -> hall -> home, OR
2. start -> home
If an amphipod is home and everything below it is home, it won’t move

With those rules in hand, we can start iterating on possible next states! From the starting state, there are exactly 28 possible moves:

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

Each of the 4 amphipods at the top of their houses (B, C, B, and D) can move to any of the valid hallway spots. Each of the ones on the bottom (A, D, C, and A) can’t move. The next function we need to write looks at each amphipod, determines if it can move, and if so, generates the next state (and its cost). Let’s get to it:

class State:
    ...

    def can_move(self, horiz: int, vert: int) -> bool:
        """
        Returns `False` if the amph is under another, `True` otherwise
        """
        if vert < 2:
            return True

        return all(self.populated.get((horiz, i)) is None for i in range(1, vert))

    def next_states(self) -> List[Tuple[int, "State"]]:
        results: List[Tuple[int, State]] = []

        for loc, amph in self.populated.items():
            horiz, vert = loc
            if not self.can_move(horiz, vert):
                continue

            if self.loc_is_home(loc) and all(
                self.loc_is_home((horiz, i))
                for i in range(vert + 1, self.max_room_size + 1)
            ):
                # me and everyone below me are home, am all good
                continue

We consider each amphipod. If it can’t move (because it’s under another), then we can bail early. Next, we check if it needs to move at all (maybe it’s already found home). If neither of those conditions match, then it’s at least possible it can move. Now we build a list of valid states for that amphipod.

The first case we consider is if it’s in a house (which is either not it’s home, or it’s blocking someone who isn’t home). We know it’s in a house because vert > 0.

from operator import add, sub

HOMES = {2, 4, 6, 8}

class State:
    ...

    def _check_horizontals(
        self, loc: GridPoint, direction: Callable[[int, int], int]
    ) -> Tuple[bool, List[Tuple[int, "State"]]]:
        """
        returns a 2-tuple of:
            *  whether or not the bath to home is unblocked (but nothing about the home itself)
            * the next possible states
        """
        results: List[Tuple[int, "State"]] = []
        horiz = loc[0]
        home_unblocked: bool = False
        while 1 <= horiz <= 9:
            horiz = direction(horiz, 1)

            if horiz == HOME_HORIZ[self.populated[loc]]:
                home_unblocked = True
                continue

            if horiz in HOMES:
                # already found home, this isn't it
                continue

            if (horiz, 0) in self.populated:
                break
            results.append(self._new_state(loc, (horiz, 0)))

        return home_unblocked, results

    def next_states(...):
        ...

        for loc, amph in self.populated.items():
            ...

            # if it's None, then we didn't check the hallway at all
            can_reach_home: Optional[bool] = None

            # so we only store the next states if we didn't go straight home
            amph_results: List[State] = []

            # if they're in a house (and thus have vert), they can only move to a hallway spot,
            # so check each direction until we hit a wall
            # this only runs if they're in a house that's not theirs
            if vert:
                for direction in [sub, add]:
                    found_home, new_states = self._check_horizontals(loc, direction)
                    can_reach_home = can_reach_home or found_home
                    amph_results += new_states

Let’s break this down. We know we can move into the hallway, but we’re not sure how far. So we call _check_horizontals to walk in a direction. If we’re on top of a house, we keep going (can’t block an entrance). If that house is ours, we note the fact that we’re not blocked from getting there (this is relevant later). For each valid spot we find, we store it as a possible state. Once we’re blocked or we reach the end of the hallway, we return all the new states we found. We can reuse the same function for both directions by passing the add and sub functions from the operators package (which are the functions that power + and -).

Back in the main loop, we check horizontally in both directions and note whether or not we could have reached home’s entrance unblocked (something we’ll need a little later). We will have found some number of new states. In our original example, this will be 7.

Next, we look at locations without vert, which are amphipods that are already in the hallway. Their only valid move is to go home. That’s only possible if the way is unblocked and there are no non-residents in the target house. We may have already found whether or not we’re unblocked. If not, we can check it now.

class State:
    ...

    def _is_horiz_clear_to_home(self, horiz: int, targ_horiz: int) -> bool:
        # haven't checked, try it
        op = add if targ_horiz > horiz else sub
        can_reach_home = True
        while horiz != targ_horiz:
            horiz = op(horiz, 1)
            if (horiz, 0) in self.populated:
                can_reach_home = False
                break
        return can_reach_home

    def next_states():
        ...

        for loc, amph in self.populated.items():
            ...

            # they're already in the hallway, so their only valid move is to go home
            # to go home, they have to be:
            # * unblocked to get there
            # * if there are other amphs there, they must all be residents

            targ_horiz = HOME_HORIZ[amph]

            if can_reach_home is None:
                can_reach_home = self._is_horiz_clear_to_home(horiz, targ_horiz)

            if not can_reach_home:
                results += amph_results
                continue

If we can’t reach home, we can save what we found so far and bail. The more interesting case is if we can. We have a couple more checks to do, before we know that it can actually land at home (not just reach it):

class State:
    ...

    def _check_vertical(
        self, loc: GridPoint, targ_horiz: int
    ) -> Optional[Tuple[int, "State"]]:
        # there is a clear path from loc to the opening of home
        # to enter home:
        # * all current residents must also be home
        # * there must be an empty spot

        targ_vert = None
        for vert_possibility in range(self.max_room_size, 0, -1):
            p = (targ_horiz, vert_possibility)
            if p in self.populated:
                if self.loc_is_home(p):
                    # can't go here, but we're not busted yet
                    continue
                # non-neighbor here!
                break

            # found an empty spot!
            targ_vert = vert_possibility
            break

        if targ_vert is None:
            # couldn't find a spot
            return None

        # otherwise, we have found home
        return self._new_state(loc, (targ_horiz, targ_vert))

    def next_states():
        ...

            if home_state := self._check_vertical(loc, targ_horiz):
                    results.append(home_state)
            else:
                results += amph_results

        return results

This covers the rules we discussed. It checks each spot in the vertical from the bottom up (starting at whatever self.max_room_size is; 2 initially). There are a couple of cases that _check_vertical handles- they’re commented. If there’s a spot and it’s not otherwise breaking rules, we can go there!

Back in the main loop, there’s some new syntax: :=. Affectionately called the “walrus operator”, it’s a combination of an if and an assignment. It’s equivalent to:

home_state = self._check_vertical(loc, targ_horiz)
if home_state:
    results.append(home_state)
else:
    results += amph_results

It’s a nice little piece of syntactic sugar. You can read more about it in the PEP that introduced it.

There’s a nice little optimization there- if we can successfully go home, we can discard the rest of the amph_results- the best move is always to go home, so we only need to bring that state along with it if it’s an option. Either way, we return whatever we have.

Congratulations on making it this far! We’ve done the lion’s share of the work- there’s just a couple of things left. Let’s revisit our Dijkstra’s implementation. Here’s what I copied from day 15 with a couple of updates for today:

def _dijkstra(self, max_room_size: int) -> int:
    locations = self._parse_input()

    start = State(locations, max_room_size)
    queue: List[Tuple[int, State]] = [(0, start)]
    visited = set()
    distances: DefaultDict[str, float] = defaultdict(lambda: inf, {start.frozen: 0})

    while queue:
        cost, current = heappop(queue)

        if current in visited:
            continue

        if current.did_win:
            return cost

        visited.add(current)

        for next_move_cost, next_state in current.next_states():
            if next_state in visited:
                continue

            total_cost = cost + next_move_cost

            if total_cost < distances[next_state]:
                distances[next_state] = total_cost
                heappush(queue, (total_cost, next_state))

    raise RuntimeError("No solution found")

If we run this as is, we get an odd error:

    ...
    if current in visited:
TypeError: unhashable type: 'State'

It all has to do with what kinds of items are allowed to be elements in sets. We’ve talked before about how Python sets are fast. They do this by taking the hash of an object- a unique string that’s determined by the object’s properties. The same object always has the same hash. Python knows very quickly if a given hash exists in a set (there’s more to it, but it’s a topic for another time).

For this approach to work, the items in the set have to be immutable, or unchangeable. Some Python data types are immutable by design: strings, numbers, etc. Containers, on the other hand, are mutable: dict, list, set. You can add and remove items to/from them, so they can change over time. As a result, you can’t put a mutable object into a place where only immutable objects are accepted: sets and dict keys. Here are some examples:

s = set()
s.add(1) # ok
s.add('a') # ok
s.add([0]) # unhashable type: 'list'

d = {}
d[1] = 1
d['a'] = 'a'
d[set()] = None # unhashable type: 'set'

This is the error we’re seeing! We’re trying to check if current (an instance of our State dataclass) is in a set. Do that, it needs to be hashable. Dataclasses can be hashable, but all of their properties must also be hashable. Since we’re storing a dict, that won’t work. Instead, we’ll have to devise a way to represent our State in a hashable way, a process known as serialization. There are a few ways to do this, but the simplest is as a string. We have to be careful to ensure that our populated dict has consistent serialized output. That requires sorting our keys, so that the same dict always has the same output, an important feature here. Here’s how that looks:

from functools import cached_property

class State:
    ...

    @cached_property
    def frozen(self) -> str:
        """
        Used to serialize this state into a set
        """
        return "|".join(f"{k}:{v}" for k, v in sorted(self.populated.items()))

Perfect! This let’s us represent our state as a hashable string. We used cached_property because we’ll be using this a lot and we don’t want to have to re-calculate it multiple times. Here’s the updated Dijkstra:

def _dijkstra(self, max_room_size: int) -> int:
    locations = self._parse_input()

    start = State(locations, max_room_size)
    queue: List[Tuple[int, State]] = [(0, start)]
    visited: Set[str] = set()
    distances: DefaultDict[str, float] = defaultdict(lambda: inf, {start.frozen: 0})

    while queue:
        cost, current = heappop(queue)

        if current.frozen in visited:
            continue

        if current.did_win:
            return cost

        visited.add(current.frozen)

        for next_move_cost, next_state in current.next_states():
            if next_state.frozen in visited:
                continue

            total_cost = cost + next_move_cost

            if total_cost < distances[next_state.frozen]:
                distances[next_state.frozen] = total_cost
                heappush(queue, (total_cost, next_state))

    raise RuntimeError("No solution found")

That should be everything! Give it a run and…

...
    heappush(queue, (total_cost, next_state))
TypeError: '<' not supported between instances of 'State' and 'State'

Ah! One last thing. A priority queue sorts its items, but it doesn’t know how to sort State instances. That makes sense - they don’t really have a defined order. All that remains is to tell Python as much:

class State:
    ...

    def __lt__(self, other):
        # needed so States can be in sortable tuples; we don't actually care about order
        return False

The __lt__ builtin is how you tell Python that an instance (self) is less than another object. And again, we don’t care about the actual result, just that Python can call this method and get an “answer”.

Finally, we’re all set. We have an answer!

Part 2

I cheated a bit on my writeup today- my part 1 answer includes everything we needed for part 2. I figured that was easier than introducing all my code that assumed there were only room sizes of 2. So, we can use our adjustable room sizes to change how we parse input:

def _parse_input(self, extra: bool) -> Dict[GridPoint, Amphipod]:
    start_order: List[Amphipod] = re.findall(r"[ABCD]", self.input)
    locations = {}
    i = 0
    # the amphs we read from input are farther down if we're in "extra" mode
    for vert in [1, 4 if extra else 2]:
        for horiz in [2, 4, 6, 8]:
            locations[(horiz, vert)] = start_order[i]
            i += 1
    if extra:
        # D#C#B#A #
        # D#B#A#C #
        # we can just hardcode these
        locations.update(
            {
                (2, 2): "D",
                (2, 3): "D",
                (4, 2): "C",
                (4, 3): "B",
                (6, 2): "B",
                (6, 3): "A",
                (8, 2): "A",
                (8, 3): "C",
            }
        )
    return locations

def _dijkstra(self, max_room_size: int) -> int:
    locations = self._parse_input(max_room_size != 2)
    ...

Thanks to our prep, our two parts become calls to that function:

def part_1(self) -> int:
    return self._dijkstra(2)

def part_2(self) -> int:
    return self._dijkstra(4)

Both parts run for me in under 5 seconds, which seems like a reasonable place to leave it. Great job sticking through with this one! I was frustrated with it for a while, but it definitely grew on me.