87 lines
2.9 KiB
Haskell
87 lines
2.9 KiB
Haskell
module Main where
|
|
|
|
import Data.Char (digitToInt)
|
|
import Data.Either (rights)
|
|
import Data.List (foldl')
|
|
import Data.Map (Map)
|
|
import qualified Data.Map as M
|
|
import System.Environment (getArgs)
|
|
|
|
-- | A candidate is either a left-most number we've just selected, or
|
|
-- a complete candidate represented as a pair.
|
|
type Candidate = Either Int (Int, Int)
|
|
|
|
main :: IO ()
|
|
main = do
|
|
input <- readFile . head =<< getArgs
|
|
putStrLn $ "Part 1: " ++ show (part1 input)
|
|
putStrLn $ "Part 2: " ++ show (part2 input)
|
|
|
|
-- | Maintain a list of candidates (either partial, or complete) and
|
|
-- then select the greatest one for each line.
|
|
--
|
|
-- Insight: By keeping a list of possible candidates we can always
|
|
-- backtrack out of an accidental corner (like making a partial
|
|
-- candidate of the last digit).
|
|
part1 :: String -> Int
|
|
part1 = sum . map (maximum . map combine . rights . foldl go [] . map digitToInt) . lines
|
|
where
|
|
go :: [Candidate] -> Int -> [Candidate]
|
|
go [] x = pure $ Left x
|
|
go (Left a : candidates') x
|
|
| x > a = Left x : Right (a, x) : candidates'
|
|
| otherwise = Right (a, x) : candidates'
|
|
go candidates@(Right (a, b) : candidates') x
|
|
| x > a = Left x : Right (a, x) : candidates
|
|
| x > b = Right (a, x) : candidates'
|
|
| otherwise = candidates
|
|
|
|
combine :: (Int, Int) -> Int
|
|
combine (a, b) = a * 10 + b
|
|
|
|
type CandidateN = [Int]
|
|
|
|
-- | Maintain one partial candidates for each length.
|
|
--
|
|
-- Insight: If there are at least N-1 digits left after we've spotted
|
|
-- the greatest digit, that digit needs to be the first digit of our
|
|
-- candidates.
|
|
--
|
|
-- Insight 2: We only ever need to store one candidate per length.
|
|
part2 :: String -> Int
|
|
part2 = sum . map (combine . (M.! 12) . foldl' checkPotential mempty . map digitToInt) . lines
|
|
where
|
|
checkPotential :: Map Int CandidateN -> Int -> Map Int CandidateN
|
|
checkPotential cs x =
|
|
let partials = filter ((< 12) . length) $ M.elems cs
|
|
partials' = [x] : map (genPartial x) partials
|
|
in foldl' insertMaxCandidate cs $ filter ((<= 12) . length) partials'
|
|
|
|
-- generates new partials by appending x at candidate or after
|
|
-- replacing a shedding smaller digits.
|
|
genPartial :: Int -> CandidateN -> CandidateN
|
|
genPartial x candidate = candidate ++ [x]
|
|
|
|
-- replaces a candidate with a better one with the same length
|
|
insertMaxCandidate :: Map Int CandidateN -> CandidateN -> Map Int CandidateN
|
|
insertMaxCandidate m c = M.insertWith maxCandidate (length c) c m
|
|
|
|
maxCandidate :: CandidateN -> CandidateN -> CandidateN
|
|
maxCandidate a b
|
|
| combine a > combine b = a
|
|
| otherwise = b
|
|
|
|
combine :: CandidateN -> Int
|
|
combine xs = sum $ zipWith power [0 :: Int ..] (reverse xs)
|
|
|
|
power :: Int -> Int -> Int
|
|
power pow b = b * 10 ^ pow
|
|
|
|
testInput :: String
|
|
testInput =
|
|
unlines
|
|
[ "987654321111111",
|
|
"811111111111119",
|
|
"234234234234278",
|
|
"818181911112111"
|
|
]
|