95 lines
3 KiB
Haskell
95 lines
3 KiB
Haskell
-- requires cabal install --lib megaparsec parser-combinators heap vector
|
|
module Main (main) where
|
|
|
|
import Control.Monad (void, when)
|
|
import Data.Functor
|
|
import qualified Data.List as L
|
|
import qualified Data.Map as M
|
|
import Data.Maybe
|
|
import Data.Ord (comparing)
|
|
import Data.Void (Void)
|
|
import Text.Megaparsec
|
|
import Text.Megaparsec.Char
|
|
|
|
import Debug.Trace
|
|
|
|
type Input = [String]
|
|
|
|
type Parser = Parsec Void String
|
|
|
|
parseInput' :: Parser Input
|
|
parseInput' = some (some alphaNumChar <* eol) <* eof
|
|
|
|
parseInput :: String -> IO Input
|
|
parseInput filename = do
|
|
input <- readFile filename
|
|
case runParser parseInput' filename input of
|
|
Left bundle -> error $ errorBundlePretty bundle
|
|
Right input' -> return input'
|
|
|
|
data Key = U | D | L | R deriving (Eq)
|
|
instance Ord Key where
|
|
compare R _ = LT
|
|
compare U _ = LT
|
|
compare D _ = LT
|
|
compare _ _ = LT
|
|
type Coord = (Int, Int)
|
|
type Memo = M.Map (Int, Char, Char) Int
|
|
|
|
keyPad :: Char -> Coord
|
|
keyPad '7' = (0, 0)
|
|
keyPad '8' = (1, 0)
|
|
keyPad '9' = (2, 0)
|
|
keyPad '4' = (0, 1)
|
|
keyPad '5' = (1, 1)
|
|
keyPad '6' = (2, 1)
|
|
keyPad '1' = (0, 2)
|
|
keyPad '2' = (1, 2)
|
|
keyPad '3' = (2, 2)
|
|
keyPad '0' = (1, 3)
|
|
keyPad 'A' = (2, 3)
|
|
keyPad '^' = (1, 0)
|
|
keyPad 'B' = (2, 0)
|
|
keyPad '<' = (0, 1)
|
|
keyPad 'v' = (1, 1)
|
|
keyPad '>' = (2, 1)
|
|
|
|
paths :: (Char, Char) -> [String]
|
|
paths (a, b) | a == '<' || x1 == 0 && y2 == 3 = [two]
|
|
| b == '<' || y1 == 3 && x2 == 0 = [one]
|
|
| otherwise = L.nub [one, two]
|
|
where
|
|
one = moves ++ "B"
|
|
two = reverse moves ++ "B"
|
|
moves = replicate (abs $ y2 - y1) (if y2 > y1 then 'v' else '^') ++ replicate (abs $ x2 - x1) (if x2 > x1 then '>' else '<')
|
|
(x1, y1) = keyPad a
|
|
(x2, y2) = keyPad b
|
|
|
|
compute :: Input -> Int
|
|
compute codes = sum $ map complexity $ zip (map (computeOne 25) codes) codes
|
|
where
|
|
complexity :: (Int, String) -> Int
|
|
complexity (output, code) = output * (read $ init code)
|
|
computeOne :: Int -> String -> Int
|
|
computeOne n = fst . computeOne' M.empty (n + 1) 'A'
|
|
where
|
|
computeOne' :: Memo -> Int -> Char -> String -> (Int, Memo)
|
|
computeOne' memo 0 _ s = (length s, memo)
|
|
computeOne' memo n start s = L.foldl' aggregate (0, memo) $ zip (start : s) s
|
|
where
|
|
aggregate :: (Int, Memo) -> (Char, Char) -> (Int, Memo)
|
|
aggregate (total, memo) (a, b) = case (n, a, b) `M.lookup` memo of
|
|
Just l -> (total + l, memo)
|
|
Nothing ->
|
|
let (l, memo') = case paths (a, b) of
|
|
[one] -> computeOne' memo (n - 1) 'B' one
|
|
[one, two] ->
|
|
let (l1, memo'') = computeOne' memo (n - 1) 'B' one
|
|
(l2, memo''') = computeOne' memo'' (n - 1) 'B' two
|
|
in if l1 < l2 then (l1, memo''') else (l2, memo''')
|
|
in (total + l, M.insert (n, a, b) l memo')
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- parseInput "input"
|
|
print $ compute input
|