Day21.hs

-- |
-- Module:         Day21
-- Description:    <https://adventofcode.com/2024/day/21 Day 21: Keypad Conundrum>
module Day21 (solve) where

import Control.Exception (assert)
import Data.Array.Unboxed (UArray, listArray, range, (!))
import Data.Char (digitToInt, isDigit)
import Data.Function (on)
import Data.List (permutations)
import Data.Text (Text)
import Data.Text qualified as T (foldl', lines, unpack)

luts :: [UArray ((Int, Int), (Int, Int)) Int]
luts =
  iterate buildLut . listArray arrayBounds $
    [abs (x2 - x1) + abs (y2 - y1) + 1 | ((x1, y1), (x2, y2)) <- range arrayBounds]
  where
    (minXY, maxXY) = ((0, -1), (2, 3))
    arrayBounds = ((minXY, minXY), (maxXY, maxXY))
    buildLut lut =
      listArray arrayBounds $
        [ foldl' min maxBound $ map cost $ filter check $ permutations moves
        | ((x1, y1), (x2, y2)) <- range arrayBounds,
          let moves =
                replicate (x2 - x1) (2, -1)
                  ++ replicate (y2 - y1) (1, 0)
                  ++ replicate (y1 - y2) (1, -1)
                  ++ replicate (x1 - x2) (0, -1)
              check = notElem (0, 0) . scanl move (x1, y1)
              cost string = sum $ zipWith (curry (lut !)) ((2, 0) : string) (string ++ [(2, 0)])
        ]
    move (x, y) (1, 0) = (x, y + 1)
    move (x, y) (0, -1) = (x - 1, y)
    move (x, y) (1, -1) = (x, y - 1)
    move (x, y) (2, -1) = (x + 1, y)
    move pos _ = assert False pos

solve :: Int -> Text -> Int
solve depth input = sum [cost (T.unpack line) * T.foldl' accumDigits 0 line | line <- T.lines input]
  where
    cost string = sum $ zipWith (curry (luts !! depth !) `on` pos) ('A' : string) string
    accumDigits n d | isDigit d = 10 * n + digitToInt d
    accumDigits n _ = n
    pos '0' = (1, 0)
    pos '1' = (0, 1)
    pos '2' = (1, 1)
    pos '3' = (2, 1)
    pos '4' = (0, 2)
    pos '5' = (1, 2)
    pos '6' = (2, 2)
    pos '7' = (0, 3)
    pos '8' = (1, 3)
    pos '9' = (2, 3)
    pos _ = (2, 0)