AoC Day 21: Keypad Conundrum

Today’s Advent of Code problem, “Keypad Conundrum”, is a contender for best problem of the year. First off, not directly 2D grids-based. Then, an interesting implementation.¹ An all-too predictable part 2 with a non-trivial solution. Took me quite some time, too.

Ok, let’s get going. Imports first, this being literate Haskell.

import Control.Applicative ((<|>))
import Control.Arrow ((&&&),first)
import Control.Lens ((^.),_1,_2)
import Control.Monad (guard,void,foldM)
import Data.Array (Array,listArray,(!),inRange,bounds,assocs)
import Data.Char (isDigit)
import Data.Function (fix)
import Data.List (find,nub)
import Data.List.Split (split,endsWith)
import Data.Maybe (mapMaybe)
import Data.MemoCombinators (memo3,bool,bits,list,char)
import Linear.V2 (V2(V2))
import qualified Data.Set as Set

type V = V2 Int

The entire problem revolves around keypads. Keypads to solve the in-story problem, and keypads to program robots to do so recursively. Our representation will include its layout, and the current position the robot that’s operating it is pointed at.

data Keypad = Keypad
  { kpPos :: V
  , kpPad :: Array V (Maybe Char)
  }
  deriving (Eq,Ord)

Two types of keypads are used: the most common is the directional keypad, and the last one in the chain is a numeric keypad. All have a static layout.

dirpad :: Keypad
dirpad = Keypad dirpadpos dirpadpad where
  dirpadpad = listArray (V2 0 0,V2 1 2)
              [ Nothing,  Just '^', Just 'A'
              , Just '<', Just 'v', Just '>'
              ]
  dirpadpos = V2 0 2

numpad :: Keypad
numpad = Keypad numpadpos numpadpad where
  numpadpad = listArray (V2 0 0,V2 3 2)
              [ Just '7', Just '8', Just '9'
              , Just '4', Just '5', Just '6'
              , Just '1', Just '2', Just '3'
              , Nothing,  Just '0', Just 'A'
              ]
  numpadpos = V2 3 2

The main operation we can perform are those resulting from directional keypad buttons: activate or move. Our result type is Maybe to be able to express forbidden moves as Nothing, and returns both the keypad’s new state and potentially an outgoing action.

Invalid positions are refused at move time, so activation always works.

kpFollow :: Keypad -> Char -> Maybe (Keypad,Maybe Char)
kpFollow kp 'A' = do
  let Just m = kpPad kp ! kpPos kp
  Just (kp,Just m)
kpFollow kp d = do
  let v = case d of
        '^' -> V2 (-1) 0
        '<' -> V2 0 (-1)
        'v' -> V2 1 0
        '>' -> V2 0 1
      p' = kpPos kp + v
  guard (inRange (bounds (kpPad kp)) p')
  void (kpPad kp ! p')  -- refuse the empty gap
  Just (kp { kpPos = p' },Nothing)

We’re going to want to combine those with the lock. Let’s define a class for all of those to follow.

class Ord s => StateMachine s where
  stateGoal :: s -> Bool
  stateFollow :: s -> Char -> Maybe (s,Maybe Char)

instance StateMachine Keypad where
  stateGoal = const False
  stateFollow = kpFollow

This allows to combine them, resulting in a combination that follows the same interface.

instance StateMachine s => StateMachine [s] where
  stateGoal [s] = stateGoal s
  stateGoal (_:ss) = stateGoal ss
  stateFollow (s:ss) i = do
    (s',mbM) <- stateFollow s i
    (ss',m) <- maybe (Just (ss,Nothing)) (stateFollow ss) mbM
    Just (s':ss',m)

With the current definitions, our goal flag is always going to be False, either directly from the keypads or transitively from composition. Success is attained when the numeric keypad completes the appropriate sequence. We can implement this as another StateMachine.

newtype Sequence = Sequence String deriving (Eq,Ord)
instance StateMachine Sequence where
  stateGoal (Sequence s) = null s
  stateFollow (Sequence s) k = case s of
    [] -> Nothing
    (h:t) | k == h -> Just ((Sequence t),Nothing)
          | otherwise -> Nothing

Oops, this makes my combination operator work on a heterogeneous list, this won’t do. Adding a wrapper type and a bit of trivial plumbing.

data SM = SMKeypad Keypad | SMSequence Sequence deriving (Eq,Ord)
instance StateMachine SM where
  stateGoal (SMKeypad kp) = stateGoal kp
  stateGoal (SMSequence s) = stateGoal s
  stateFollow (SMKeypad kp) i = first SMKeypad <$> stateFollow kp i
  stateFollow (SMSequence s) i = first SMSequence <$> stateFollow s i

Yay. We’ve got all we need to reify a state for the entire chain of robots!

chain :: Int -> String -> [SM]
chain depth code =
  replicate depth (SMKeypad dirpad) ++
  [SMKeypad numpad,SMSequence (Sequence code)]

We can now implement a good old BFS and search it!

bfs :: StateMachine s => [Char] -> s -> Int
bfs is s0 = go Set.empty [(s0,0)] where
  go cl ((p,d):q)
    | p `Set.member` cl = go cl q
    | stateGoal p = d
    | otherwise = go cl' (q ++ q') where
        cl' = Set.insert p cl
        q' = (,d+1) . fst <$> mapMaybe (stateFollow p) is

And this solves part 1.

part1 :: [String] -> Int
part1 = sum . map complexity where
  complexity code =
    bfs "^<v>A" (chain 2 code) * read (takeWhile isDigit code)

Part 2, what a surprise, extends the chain. A lot. Full search doesn’t cut it anymore. Back to the drawing board.

There’s a few things to notice to make it all faster.

First, full search on a single level is usually kind of wasteful. The target keypad is not a maze, you can always go from a point to the other by following a simple path of two segments, one horizontal and one vertical. The controlling robot has to press A for each move, so it’s obviously shorter overall not to move unless the current direction is exhausted.

Let’s write a function to return the one or two reasonable paths from a position to the next. In the list monad, we’ll make deeper use of it shortly.

kpPathToPress :: Keypad -> Char -> Char -> [String]
kpPathToPress kp src dst = map (++ "A") $ nub $
  let p1 = kpFind kp dst
      p2 = kpFind kp src
      V2 v h = p1 - p2
      vPath = replicate (abs v) ("^ v" !! (signum v + 1))
      hPath = replicate (abs h) ("< >" !! (signum h + 1))
      hole = kpFindHole kp
      notSame coord a b = (b - a) ^. coord /= 0
  in (vPath ++ hPath) <$ guard (notSame _1 p1 hole || notSame _2 p2 hole) <|>
     (hPath ++ vPath) <$ guard (notSame _2 p1 hole || notSame _1 p2 hole)

This makes use of a bit of trivial support code:

kpFind :: Keypad -> Char -> V
kpFind kp c = kpFindRaw kp (Just c)

kpFindHole :: Keypad -> V
kpFindHole kp = kpFindRaw kp Nothing

kpFindRaw :: Keypad -> Maybe Char -> V
kpFindRaw kp c = p where
  Just p = fst <$> find ((== c) . snd) (assocs (kpPad kp))

I found the other relevant observation nicely hinted in the problem statement. Notice how it says you have to type multiple codes on the numeric keypads, but asks their complexity independently? How does this make sense?

It turns out it doesn’t make a difference whether they’re done separately, in sequence, in reverse order or anything. Why? Because the robots always reach a keypad pointing at the A button. And they always complete a (shortest) sequence—at their level—on the A button again. Plus, the code sequences end on As themselves.

This makes the problem entirely decomposable into shorter subproblems, whose results can be cached.

Let’s outline a solve function, with recursion abstracted away.

solve :: (Bool -> Int -> String -> Int) -> Bool -> Int -> String -> Int

solve rec numOrDir depth seq is intended to return the length of a shortest sequence by the human operator to have a keypad, type decided by numOrDir, isolated behind depth layers of directional keypads, generate the key sequence seq.

For a depth of 0, the human is typing the sequence directly, that’s easy enough.

solve _ _ 0 s = length s

For deeper keypads, we decompose the sequence into atomic cacheable chunks, generate the list of reasonable robotic arm paths that can yield it, keep the length of the smallest chain of human actions that yield it, and sum.

That length of the smallest chain of human actions conveniently happens to be a solve result one level shallower.

solve rec isNum depth s =
  sum $
  map (
    minimum .
    map (rec False (depth - 1)) .
    kpSeq (if isNum then numpad else dirpad)
  ) $
  split (endsWith "A") s

A bit of support code to generate those lists of sequences. In the worst case, there’s two potential paths to go from a key to the next (horizontal first or vertical first), so the typical recursion patterns² will be like v<<A, two choices of ordering for a sequence of meaningful length 3 since repeated presses are just more A at the level above. Which is to say, the combinatorics here are more than tractable, they’re small.

kpSeq :: Keypad -> [Char] -> [String]
kpSeq kp s =
  foldM (\a bs -> (a ++) <$> bs) [] $
  zipWith (kpPathToPress kp) ('A':s) s

Sprinkle a bit of memoization, and we’re ready to solve part 2.

part2 :: [String] -> Int
part2 = sum . map complexity where
  complexity code =
    fix (memo3 bool bits (list char) . solve) True 26 code
    * read (takeWhile isDigit code)

A wrapper to make sure we keep the bad-performance part 1 solution in scope, and we’re done.

main :: IO ()
main = interact $ show . (part1 &&& part2) . lines

A puzzle without a dedicated parse function. How refreshing too!

This concludes today’s code write-up. See you tomorrow!

This problem would have deserved a better solution than what I present here. Better, stricter typing, less static data hauled around. But I have real-life constraints too, and don’t want to get to solving in more than 24 hours, so at some point I dropped the niceties and just shoved Chars around. A pity. And I probably won’t ever get back to it. Oh well.↩︎
The utmost typical one is just A.↩︎