Advent of Code 2024 day 21

Day 21 was, I think, the hardest problem so far. Not so much the complexity of the final code, but the complexity of the thinking that led to it.

Data types

I started by writing done the types that would be important. The robots can do actions, and there are two types of button: numeric and action. That implies the buttons should be a typeclass. Each of the button types has three functions defined for it: where is a button, what is the A button, and what positions are legal on this keypad.

type Position = V2 Int -- r, c

data Action = R | U | D | L | A deriving (Eq, Ord, Show, Enum, Bounded)
type ActionSeq = [Action]

class Button a where
  buttonPos :: a -> Position
  aButton :: a
  legalPos :: a -> Position -> Bool

instance Button Action where
  buttonPos U = V2 0 1
  buttonPos A = V2 0 2
  buttonPos L = V2 1 0
  buttonPos D = V2 1 1
  buttonPos R = V2 1 2

  aButton = A

  legalPos _ p = p `elem` [V2 0 1, V2 0 2, V2 1 0, V2 1 1, V2 1 2]

instance Button Char where
  buttonPos '7' = V2 0 0
  buttonPos '8' = V2 0 1
  ...
  buttonPos 'A' = V2 3 2

  aButton = 'A'

  legalPos _ p = p `elem` [V2 0 0, V2 0 1, ... V2 3 2]

Part 1

This is direct simulation of the various robots. I knew it wasn't going to be the right approach for part 2, but I didn't want to make the wrong guess for what part 2 could be, and waste some effort.

First try

I started off simple: the moves for a code are found by finding the movesBetween for each sequential pair of buttons, always starting a A.

moves :: Button a => [a] -> ActionSeq
moves bs = concatMap moveBetween $ zip (aButton : bs) bs

moveBetween :: Button a => (a, a) -> ActionSeq
moveBetween (a, b) = sort (mh ++ mv) ++ [A]
  where aPos = buttonPos a
        bPos = buttonPos b 
        V2 dr dc = bPos ^-^ aPos
        mh = replicate (abs dc) (if dc > 0 then R else L)
        mv = replicate (abs dr) (if dr > 0 then D else U)

I rely in the sorting order of Action to ensure the robot arm always remained in a legal position: move right first, then up and down, then finally left.

I handled the sequence of robots by chaining calls to moves.

complexity :: String -> Int
complexity code = (length ms) * ns
  where ms = moves $ moves $ moves code
        ns = read $ filter isDigit code

That worked for most of the examples, but gave the wrong answer for the last one. I was missing that different orderings of moves could lead to different lengths of expanded sequences.

Second try

The choices meant the moveBetween had to become more complex. Rather than returning the one sequence of moves between a and b, it had to return as list of them. Given a bunch of moves like "two up and two left", I need to find the induced costs of each permutation of those moves. Rather than checking all 24 permutations of four items, I reduce the number in in three stages:

remove the duplicate permutations, allowing for repeated moves
check that all intermediate positions are legal
check that like moves are grouped together (the sequence "up, left, up" can never be shorter than either "up, up, left" as the second "up" can be activated with no additional movement of the arm)

That last one in particular reduces the runtime from a minute and a half to about a tenth of a second!

moveBetween :: Button a => (a, a) -> [ActionSeq]
moveBetween (a, b) = filter (allLegal a) $ filter groupTogether possibles
  where aPos = buttonPos a
        bPos = buttonPos b 
        V2 dr dc = bPos ^-^ aPos
        mh = replicate (abs dc) (if dc > 0 then R else L)
        mv = replicate (abs dr) (if dr > 0 then D else U)
        possibles = fmap (++ [A]) $ nub $ permutations $ mh ++ mv
        groupTogether p = sort (group p) == group (sort p)
        allLegal a t = all (legalPos a) (positionsOf a t)

With moveBetween returning a list of action sequences rather than just one, moves also needed to change to pick options from the choices. That fits the List instance of Traversable, so picking the options is handled by sequenceA (or sequence).

moves :: Button a => [a] -> [ActionSeq]
moves bs = fmap concat $ sequence $ fmap moveBetween $ zip (aButton : bs) bs

Finally, complexity changes to find the shortest length.

complexity :: String -> Int
complexity code = (length ms) * ns
  where ms = head $ sortOn length $ concatMap moves $ concatMap moves $ moves code
        ns = read $ filter isDigit code

Part 2

This was the real brain-burner part of the puzzle for me.

The idea was simple enough. The length of each action sequence grows by a factor of 2-4 each step. It was going to be infeasible to generate the entire move sequence for each code across two dozen robots.

I opted for a dynamic programming approach. For each robot, starting with the one closest to me, I would find the costs of all possible pairs of button presses for the next robot. I would then use that to find the costs of the next robot along, and so on until I found all the robots.

The cache was a Map that gave the cost of each move at each level.

data CacheKey = CacheKey { moveFrom :: Action, moveTo :: Action, layer :: Int } 
  deriving (Eq, Ord, Show)
type Cache = M.Map CacheKey Int

The initial cache was just the distance between each pair of buttons on the robot keypad (plus 1 for the A press after each one).

cache1 :: Cache
cache1 = M.fromList $ fmap go allPairs
  where allPairs = [(a, b) | a <- [R .. A], b <- [R .. A]]
        go (a, b) = (CacheKey a b 1, 1 + mdist a b)
        mdist p q = let V2 dr dc = (buttonPos q) ^-^ (buttonPos p) 
                    in abs dr + abs dc

I can take an existing cache and extend it for one level by doing something similar: for each pair of buttons, find the cheapest cost move between them (at the level below) and add it to the cache. A repeated fold over those builds the whole cache.

buildCache :: Int -> Cache
buildCache maxLevel = foldl' extendCache cache1 [2..maxLevel]

extendCache :: Cache -> Int -> Cache
extendCache cache level = foldl' go cache allPairs
  where
    allPairs = [(a, b) | a <- [R .. A], b <- [R .. A]]
    go c (a, b) = M.insert (CacheKey a b level) 
                           (cheapestCostMove c (level - 1) (a, b)) c

I find the cheapestCostMove by finding the action sequences for that move and looking up their costs in the existing cache (with sequenceCostUsingCache). moveCostUsingCache does the actual lookup for each move.

cheapestCostMove :: Button a => Cache -> Int -> (a, a) -> Int
cheapestCostMove cache level (a, b) = 
  minimum $ fmap (sequenceCostUsingCache cache level) stepChoices
  where stepChoices = moveBetween (a, b)

sequenceCostUsingCache :: Cache -> Int -> ActionSeq -> Int
sequenceCostUsingCache cache level bs = 
  sum $ fmap (moveCostUsingCache cache level) $ zip (aButton : bs) bs

moveCostUsingCache :: Cache -> Int -> (Action, Action) -> Int
moveCostUsingCache cache level (a, b) = 
  M.findWithDefault (maxBound :: Int) (CacheKey a b level) cache

You can tell that this took a lot of putting together from the number of very small functions I defined. I had to think in small pieces to work through the logic!

Building the cache like this gives me the cost of each robot in the chain. I still need to find the cost of typing the keypad. This still uses moves to find the possible action sequences (and their may be several of them), then sequenceCostUsingCache to find the costs instead of the chain of concatMap moves.

complexity2 :: Cache -> String -> Int
complexity2 cache code = ms * ns
  where ms = minimum $ fmap (sequenceCostUsingCache cache 25) $ moves code
        ns = read $ filter isDigit code

The whole thing runs in about a tenth of a second, including building the cache.

Code

You can get the code from my locally-hosted Git repo, or from Codeberg.