Advent of Code Day 1: 3-Sum and the Maybe Monad

Advent of Code 2020 is underway and the first problem is a Leetcode classic - 2-Sum, then 3-Sum. They’re pretty similar, so I’ll focus on just 3-Sum.

naiveAnswer :: [Int] -> [Int]
naiveAnswer inp = do
  let indexedInput = zip inp [0 ..]
  (x, ix) <- indexedInput
  (y, iy) <- indexedInput
  (z, iz) <- indexedInput
  if length [ix, iy, iz] 
        /= length (nub [ix, iy, iz]) 
        || (x + y + z) /= 2020
    then []
    else return (x * y * z)

Without regard for performance, the notation using the Maybe monad is quite nice. We want to find three numbers in the list who sum to 2020, and the <- statements in the do-notation suggest correctly that we’re just drawing arbitrary numbers from the list. (We zip indices into the list to make sure we’re not reusing the same list entry twice.) Aside from the indices, the result reads just like the problem specification.

Of course, each of those <- statements is iterating over indexedInput, so this approach is cubic in the length of the list. (Which actually turns out to be fine for the size of the given puzzle input.) Luckily, we can get down to quadratic time (which is conjectured to be asymptotically optimal) in the same number of lines:

betterAnswer :: [Int] -> [Int]
betterAnswer inp = do
  let indexedInput = zip inp [0 ..]
      numToIndices = 
        HashMap.fromListWith (<>) 
          $ map (second (: [])) indexedInput
  (x, ix) <- indexedInput
  (y, iy) <- indexedInput
  let z = 2020 - x - y
      zInds = 
        HashMap.lookupDefault [] z numToIndices
  if any (\iz -> iz /= ix && iz /= iy) zInds 
    then return (x * y * z) 
    else []

A lot of the terseness here is down to the niceness of having a function like HashMap.fromListWith available. (What about (: [])? That’s the (:) operator, partially-applied-on-the-right to [] – a funny way of spelling the singleton function.) We construct a multimap from number values to the indices at which they appear in the list. Then we draw arbitrary values x and y from the list and check to see if a z such that x + y + z == 2020 is in that multimap.

My complete runnable solution is on Github: https://github.com/ewilden/aoc-2020/blob/main/src/Day01.hs