Did you ever play a snake game? It was one of the first things I tried to program as a kid, using character graphics mind you, writing Z80 assembly language. This challenge is about simulating the snake's path given the list of movements, while keeping track of the tail's position.

I'll say the tricky part is correctly «following the leader». It was a fun «aha» moment when I was a teenager. So much so, that it stuck with me, and being able to use it again made my day.

The Problem

There's an infinite two-dimensional grid. There's a knotted string, that I shall refer as «the snake» from now own, sitting at the origin, i.e. row 0, column 0. The snake has several segments: Part A says there's two, and I figured Part B was going to ask for more.

The snake can move up, down, left, or right. Our input is a list of said movements, including the number of steps in the particular direction:

    R 4\n
    U 4\n
    L 3\n
    D 1\n
    R 4\n
    D 1\n
    L 5\n
    R 2\n

The «head» of the snake moves the given number of steps in the particular direction, and the trailing segments «follow» the head. I'll discuss the following rules while explaining their implementation.

We'll need to parse and simulate the movements. The problem requires keeping track of all the unique positions in the grid that the tail of the snake has visited.

Main Program

The main program is pretty much the same as Day 1.

Loading movements

Simple lazy list parsing and a couple data types suffice for handling the input

    data Dir  = L | U | R | D
              deriving (Show,Read,Eq)

    data Move = Move { direction :: Dir
                     , steps     :: Int
                     }
              deriving (Show,Eq)

Making Dir an instance of Read allows me to use read to go from the Char into the proper constructor. Other than that, processing each line out of lines, using break to get its parts, and pattern matching to make a Move, leads to

    loadMoves :: String -> [Move]
    loadMoves = map loadMove . lines
                  where
                    loadMove :: String -> Move
                    loadMove x = Move (read d) (read s)
                      where
                        (d,s) = tail <$> break (==' ') x

We know focus on modeling the snake, and simulating one Move at a time

Moving a snake

Positions and their closeness

It's better to model the position of each snake's segment, rather than go about simulating an infinite grid. Let us model a single two-dimensional position with

    data Pos  = Pos { row :: Int
                    , col :: Int
                    }
              deriving (Show,Eq,Ord)

Focusing on individual positions helps describing how to take a single step in a particular direction

    step :: Pos -> Dir -> Pos
    step (Pos r c) U = Pos (r+1)  c
    step (Pos r c) D = Pos (r-1)  c
    step (Pos r c) L = Pos  r    (c-1)
    step (Pos r c) R = Pos  r    (c+1)

and later apply step repeatedly to accomplish multiple step movements.

Now that we have Pos to model a single segment position, we turn to the problem of deciding whether or not two positions are close or not. Here's the first «math realization»: they must be no more than one row or column apart in any direction. The names of the arguments reflect that I will usually compare the «current segment» with the «next segment».

    stillClose :: Pos -> Pos -> Bool
    stillClose current next = abs (col current - col next) <= 1 &&
                              abs (row current - row next) <= 1

A collection of positions

Thinking of a Snake as an ordered collection of positions leads to

    type Snake = [Pos]

Nothing fancy. The head of the list will be the head of the Snake. The ordering of the segments is implicit. I can take advantage of all of Haskell's lazy list processing to make it move.

Since the snake always starts at the origin with all its segments one on top of the other, it made sense to write

    initialSnake :: Int -> Snake
    initialSnake n = replicate n (Pos 0 0)

to create an n-segment Snake sitting at the origin.

The moving trick

The problem description has a long written explanation, with supporting ASCII-art, showing how the «tail» follows the «head» in the two-segment case. I summed them up as a couple of rules:

1. Each segment and the one following it must always be touching, vertically, horizontally, or diagonally. For a two-segment snake it means the «tail»'s coordinates must be either the same as or any of the eight immediate neighbors of the «head»'s coordinates. If the snake has more than two segments, the same rule applies for any two consecutive segments.

2. If the «leading» segment ever moves away two steps, so that it is not adjacent to the «trailing» segment, the latter is «pulled» close again. If the «leading» segment moved horizontally (vertically), the «trailing» segment will advance in the same fashion. But if the «leading» segment moved such that the «trailing» segment is on a different row and different column, then the «trailing» segment must move diagonally immediately «behind» the «leading» segment.

The first rule makes processing an arbitrarily long Snake easier. We write

    pull :: Snake -> Dir -> Snake
    pull (m:ms) d = reverse $ foldl' go [step m d] ms
      where
        go :: [Pos] -> Pos -> [Pos]
        go chain@(prev:_) curr = if stillClose prev curr 
                                    then curr:chain
                                    else fix prev curr:chain

Move the first Pos of the Snake one step in whatever direction. We'll save our «adjusted» positions into an accumulator list. Go left to right on the rest of the Snake Positions, considering two elements at a time: the last one added to the accumulator, and the current one in the list. Check if the current one is stillClose to the previous one. Extend the chain as is if they are; if they're not, fix the current one in relation to the previous one before extending the chain.

We are using foldl' so the resulting accumulator chain needs to be reversed to have the updated Snake in the proper order.

The second rule establishes how to pull an element close to fix it. This is where the second «math realization» comes into play. Let H be the position of the «head» after moving. After using stillClose we notice that the «tail» T must be pulled close. It follows that T can only be in the places shown below

      |T| |T|
     -+-+-+-+-
     T| |X| |T
     -+-+-+-+-
      |X|H|X|
     -+-+-+-+-
     T| |X| |T
     -+-+-+-+-
      |T| |T|

needing to move to the closest X. Here's the «math realization»:

• Let (rh,ch) be the row and column coordinates for H.

• Let (rt,ct) be the row and column for any of the T above.

• ch - ct will always be -1, 0, or 1 -- the «column difference»

• rh - rt will always be -1, 0, or 1 -- the «row difference»

So, for any T position, we only need to add the column difference to its column, and the row difference to its row, thus computing the desired destination. We complete the definition for pull by adding the local function

        fix    :: Pos -> Pos -> Pos
        fix (Pos rh ch) (Pos rt ct) = Pos r' c'
          where 
            r' = rt + signum (rh - rt)
            c' = ct + signum (ch - ct)

But pull only moves the Snake one step in a particular direction...

Large scale movement

Every Move has one or more steps. We can «expand» a [Move] into a list of single step movements [Dir]: turn each Move into a [Dir] and concat the results. We can then use pull one step at a time, starting from the initialSnake of whatever size we need, and working our way through the [Dir].

Every time we apply pull on a Snake we compute the updated Snake after the movement. We are asked to figure out the position of the last segment of every Snake along the way. This is a perfect opportunity to use scanl'

    scanl' :: (b -> a -> b) -> b -> [a] -> [b]

because it works like foldl' in the sense that it goes left to right using an accumulating function, while saving all intermediate results. We can see how it's perfect by plugging pull as the first argument, with substitution a ~ Snake and b ~ Dir

    scanl' :: (Snake -> Dir -> Snake)   -- pull
           -> Snake                     -- initialSnake
           -> [Dir]                     -- expanded single step movements
           -> [Snake]                   -- All the Snakes!

Here's our trail of Snakes

    trail :: Int -> [Move] -> [Snake]
    trail n ms = scanl' pull (initialSnake n) (concatMap expandMove ms)
      where
        expandMove :: Move -> [Dir]
        expandMove (Move d l) = replicate l d

A generic solution

As usual, I wanted the solution to work for arbitrarily long Snake, because needing to know the length of things is contrary to the functional programming mindset, in the general case. That's why trail's first argument specifies the size of the initial Snake, and the remaining machinery works independently.

We can loadMoves and then use trail to compute all the Snakes. If we think of this as a collection of «snapshots» of where the snake was, we can focus on the last position of each location by using map last. That is, we reason using the particular size for Part A,

    ghci> :type loadMoves
    loadMoves :: String -> [Move]
    ghci> :type trail 2 . loadMoves
    trail 2 . loadMoves :: String -> [Snake]
    ghci> :type map last . trail 2 . loadMoves
    map last . trail 2 . loadMoves :: String -> [Pos]

We just need to remove duplicates from that [Pos], to have the unique positions the last segment visited. There's a naïve solution using Data.List.nub that does the trick: nub's time complexity is O(n²). There's a much better solution that's O(n log n) if we use Data.Set.

    import qualified Data.Set as Set

    solution :: Int -> String -> Int
    solution n s = Set.size
                 $ foldl' (flip Set.insert)
                          Set.empty
                          (map last $ trail n $ loadMoves s)

Using a Set.Set as accumulator, and starting with an Set.empty one, we process the [Move] by Set.insert'ing it into the accumulator. This gets rid of duplicates as per the definition of a Set. The API provides a O(1) way to compute Set.size.

By the way, we can use Set on account of Pos being an instance of Ord.

Solving both parts

Part A is for two-segment snakes

    partA :: String -> Int
    partA = solution 2

After getting the right result, Part B is the same but using ten-segment snakes, so

    partB :: String -> Int
    partB = solution 10

Conclusion

Breaking a large problem into smaller ones, and then composing your way up, is a crucial functional programming skill. In this case, focusing on single step moves, and single segment adjustments, were two good examples.

Generalizing problems as lazy list (Foldable) processing using map, foldl and scanl, is also a crucial functional programming skill. Realizing that pulling a Snake is a foldl of fixes over individual segments is the first stage: it doesn't matter how long the Snake is. Noting the trail of Snakes is a repeated foldl of pulls while collecting the results is more of the same: knowing scanl does that comes from spending time using Haskell's standard library instead of reinventing the wheel badly.

Choosing the efficient Data.Set over the obvious naïve nub solution only comes out of study and experience. There are cues for what's proper to use on problems: «unique whatever» hints at Set more often than not.

Lists are the poor man's set. A very expensive one at it.