You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
87 lines
2.9 KiB
87 lines
2.9 KiB
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], 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 (snd . (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], 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 (ds, _) = let ds' = ds ++ [x] in (ds', combine ds') |
|
|
|
-- 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 |
|
| snd a > snd b = a |
|
| otherwise = b |
|
|
|
combine :: [Int] -> 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" |
|
]
|
|
|