advent of code, haskell

Advent of Code 2021 day 23

The solution to day 23 was the longest program so far. At 320 lines of code, this solution is 50% longer than the next longest (day 15, another search problem) and twice as long as day 16, the longest non-search solution. The program was long because it had to contain a lot of domain knowledge about the problem setting.

Understanding the domain

Fundamentally, the problem was similar to many previous ones. The arrangement of amphipods in the burrow represents a state, and there are legal moves that transform one state to another. The problem is to search the space of legal move sequences to turn the initial state into a goal state.

I'd normally handle this by allowing each amphipod to move one space and just allowing the search tree to become deep. But in this case, the additional rules on movement meant I had to take a different approach.

The problem has a constraint that, once an amphipod stops in the hallway, it can't move again until it enters its destination room. I could have handled that by expanding the notion of a problem state to include tracking which amphipod, if any, was currently moving. Instead, I decided to pre-compute "moves" a being the whole trail from start to end, and keeping those moves as atomic.

This meant there were three types of moves:

  1. From a room to the hallway
  2. From the hallway to a room
  3. From a room directly to a room (but not back to the same room)

Including the room-room moves eliminates some duplicate search states if there are several hall spaces between the source and destination rooms. It doesn't matter which one is used to break the journey from room to room, so I eliminate that choice by including them all in one move.

The other piece of domain knowledge I included was ensuring that when an amphipod enters a room, it always moves to the furthest available space. All the spaces must be filled eventually, so it makes little sense to have unfilled spaces at the far end.

Data structures

There ended up being a lot to keep track of, so I created a bunch of data structures to keep track of everything. I also used lenses a lot to access all the different parts of these structures.

Amphipods and positions have explicit types, and the position of the amphipods is wrapped in a MoveState. I also rename the lenses for the V2 constructor, as I'm using (row, column) coordinates.

data Amphipod = A | B | C | D deriving (Show, Read, Eq, Ord, Enum)

type Coord = V2 Int -- r, c
_r :: Lens' (V2 Int) Int
_r = _x
_c :: Lens' (V2 Int) Int
_c = _y

type MoveState = M.Map Coord Amphipod

A Step is a single valid move of an amphipod, including the distance travelled and the spaces moved through; the latter is used to check there's nothing in the way of a move. Steps relates the starting position (the key) with the set of final positions.

An AppliedMove combines the positions of amphipods after a move and the Step that got there.

data Step = Step 
  { _destination :: Coord
  , _distance :: Int
  , _transits :: S.Set Coord
  , _entryRequirement :: Maybe Amphipod
  } deriving (Show, Eq, Ord)
makeLenses ''Step

type Steps = M.Map Coord (S.Set Step)

data AppliedMove = AppliedMove 
  { _afterMove :: MoveState 
  , _appliedStep :: Step
  }
  deriving (Show, Eq, Ord)
makeLenses ''AppliedMove

Finally, a Burrow holds all the possible Steps and information about where the hall and rooms are. That gets fed into the Reader monad, for use inside the search routines.

data Burrow = Burrow 
  { _possibleSteps :: Steps
  , _roomColumns :: M.Map Amphipod Int
  , _hallRow :: Int
  } deriving (Show, Eq)
makeLenses ''Burrow  

type BurrowContext = Reader Burrow

Building the burrow

mkBurrow reads the text from the input file and returns the Burrow (including all the moves) and the initial arrangement of amphipods. Mainly it's just walking over the text file, pulling out bits, and recording their positions. 

mkBurrow :: String -> (Burrow, MoveState)
mkBurrow text = (burrow, initState) -- (burrow, initState)
  where rows = lines text
        hall = mkHall (rows!!1)
        rooms = mkRooms $ drop 2 rows
        roomCols = S.map (^. _c) $ M.keysSet rooms
        hall' = S.filter ((`S.notMember` roomCols) . (^. _c)) hall 
        routes = mkRoutes hall' rooms
        roomColMap = M.fromList $ zip [A .. D] $ S.toAscList roomCols
        burrow = Burrow { _possibleSteps = routes, _roomColumns = roomColMap, _hallRow = 1}
        initState = mkInitialState rows

mkHall :: String -> S.Set Coord
mkHall text = S.fromList hallCoords
  where hallCols = filter ((/= '#') . snd) $ zip [0..] text
        hallCoords = map ((V2 1) . fst) hallCols

mkRooms :: [String] -> M.Map Coord Amphipod
mkRooms text = M.unions rooms
  where rooms = map mkRoom $ zip [2..] text

mkRoom :: (Int, String) -> M.Map Coord Amphipod
mkRoom (r, text) = M.fromList roomCoords
  where roomCols = filter ((`elem` ("ABCD." :: String)) . snd) $ zip [0..] text
        roomCoords = zip (map ((V2 r) . fst) roomCols) [A .. D]

The interesting part comes in mkRoutes. It finds all the moves of each separate type (as described above) the combines them with M.unionsWith. 

mkRoutes :: S.Set Coord -> M.Map Coord Amphipod -> Steps
mkRoutes halls rooms = M.unionsWith (S.union) [hallRoutes, roomHallRoutes, roomRoomRoutes]
  where hallRoutes = S.foldr' (mkHallRoute rooms) M.empty halls
        roomHallRoutes = S.foldr' (mkRoomHallRoute halls) M.empty (M.keysSet rooms)
        roomRoomRoutes = S.foldr' (mkRoomRoomRoute hallRow rooms) M.empty (M.keysSet rooms)
        hallRow = S.findMin $ S.map (^. _r) halls

(As a little note to self, I spent a fair bit of time running in circles updating the various Maps without using S.union to combine the multiple sets for each key!)

As an example, making the routes from the hall spaces to the room spaces, mkHallRoute, is a nested fold. First I create the routes for each hall space then, for each one, the routes from that hall space to each room space.

mkHallRoute :: M.Map Coord Amphipod -> Coord -> Steps -> Steps
mkHallRoute rooms here routes = M.foldrWithKey' (mkHallRoute1 here) routes rooms

mkHallRoute1 :: Coord -> Coord -> Amphipod -> Steps -> Steps
mkHallRoute1 here@(V2 hr hc) there@(V2 tr tc) entry routes = M.insert here (S.insert step existingRoutes) routes
  where step = Step { _destination = there
                    , _distance = (S.size transits) 
                    , _transits = transits
                    , _entryRequirement = Just entry
                    }
        cMin = min hc tc
        cMax = max hc tc
        transits = S.delete here $ S.fromList $ [V2 hr c | c <- [cMin..cMax]] ++ [V2 r tc | r <- [hr..tr]]
        existingRoutes = M.findWithDefault S.empty here routes

The other types of route creation follow a very similar pattern.

Searching

Most of the search is standard A* search. The interesting bits are the problem-specific parts.

successors finds all the legal applied moves from a particular arrangment of amphipods. It does that by, for each amphipod, finding all the legal moves of that amphipod.

successors :: MoveState -> BurrowContext (S.Set AppliedMove)
successors moveState = 
  do steps <- asks (^. possibleSteps)
     let succs = M.foldrWithKey' (legalSteps steps moveState) S.empty moveState
     return succs
     
legalSteps :: Steps -> MoveState -> Coord -> Amphipod -> S.Set AppliedMove -> S.Set AppliedMove
legalSteps steps state here amphipod acc = S.union appliedSteps acc
  where allSteps = steps ! here
        occupiedSpaces = M.keysSet state
        freeSpaces st = S.null $ S.intersection occupiedSpaces (st ^. transits)
        freeSteps = S.filter freeSpaces allSteps
        validTargetSteps = S.filter (\st -> fromMaybe amphipod (st ^. entryRequirement) == amphipod) freeSteps
        openRoomSteps = S.filter (openRoom state) validTargetSteps
        highestRowSteps = S.filter (highestRow (S.map (^. destination) openRoomSteps)) openRoomSteps
        appliedSteps = S.map (\s -> AppliedMove 
                                      { _afterMove = (applyStep state here s)
                                      , _appliedStep =  s
                                      }
                              ) highestRowSteps

legalSteps finds the legal moves by finding all the moves (allSteps = steps ! here) then filtering out the illegal ones over several stages:

  • freeSteps are steps where every spaced moved through is empty
  • validTargetSteps are those to a place this amphipod could go (the hall or its own room)
  • openRoomSteps are where the target room is "open", i.e. does not have any non-destined amphipods
  • highestRowSteps are to the furthest unoccupied space in a column

Heuristic

It's A* search, so it's guided by a heuristic. The one I use relaxes the constraints of "can't pass through other amphipods" and "only one amphipod per space in the destination room".

estimateCost :: MoveState -> BurrowContext Int
estimateCost state = 
  do rCols <- asks (^. roomColumns)
     hRow <- asks (^. hallRow)
     let amphipodCosts = M.mapWithKey (estimateACost rCols hRow) state
     return $ sum $ M.elems amphipodCosts

estimateACost :: M.Map Amphipod Int -> Int -> Coord -> Amphipod -> Int
estimateACost rCols hRow (V2 r c) amphipod = (singleStepCost amphipod) * dist
    where targetCol = rCols ! amphipod
          dist = if c == targetCol
                 then 0
                 else (r - hRow) + (abs (c - targetCol)) + 1

It may not be much, but it's enough to halve the runtime of the program.

Profiling

My first attempts at this took several minutes to run part 2. This seemed like too long a time. I profiled the code and found that 96% of the runtime was in legalSteps (which wasn't a huge surprise). But the profiler, by default, works at the level of functions, so didn't give me any more detail than that.

I manually inserted some additional call centres at all the filter stages of legalSteps, with annotations like {-# SCC freeSteps #-}, like this:

legalSteps :: Steps -> MoveState -> Coord -> Amphipod -> S.Set AppliedMove -> S.Set AppliedMove
legalSteps steps state here amphipod acc = S.union appliedSteps acc
  where allSteps = steps ! here
        freeSpaces st = S.null $ S.intersection (M.keysSet state) (st ^. transits)
        freeSteps = {-# SCC freeSteps #-} S.filter freeSpaces allSteps
        validTargetSteps = {-# SCC validTargetSteps #-} S.filter (\st -> fromMaybe amphipod (st ^. entryRequirement) == amphipod) freeSteps
        openRoomSteps = {-# SCC openRoomSteps #-} S.filter (openRoom state) validTargetSteps
        highestRowSteps = {-# SCC highestRowSteps #-} S.filter (highestRow (S.map (^. destination) openRoomSteps)) openRoomSteps
        appliedSteps = S.map (\s -> AppliedMove 
                                      { _afterMove = (applyStep state here s)
                                      , _appliedStep =  s
                                      }
                              ) highestRowSteps

That showed that most of the time was spent finding the freeSteps. On a hunch that GHC wasn't bright enough to cache the result of M.keysSet state, I pulled that out into a separate named variable:

occupiedSpaces = M.keysSet state
freeSpaces st = S.null $ S.intersection occupiedSpaces (st ^. transits)
freeSteps = S.filter freeSpaces allSteps

Just that halved the runtime.

Code

You can get the code from my locally-hosed Git repo, or from Gitlab.