aoc2020/day10.hs

import Data.List (sort, subsequences)

main :: IO ()
main = do
  raw <- readFile "day10.txt"
  let nums = (map read $ lines raw) :: [Int]
      ansA = solveA nums
      ansB = solveB nums
   in do
     putStrLn $ "day10a: " ++ (show ansA)
     putStrLn $ "day10b: " ++ (show ansB)

solveA :: [Int] -> Int
solveA nums = (filt 3 ls) * (filt 1 ls)
  where
    filt a = length . (filter (== a))
    end = 3 + (maximum nums)
    ls = (diffs . sort) (0 : end : nums)

solveB :: [Int] -> Int
solveB nums = unLookup 0 $ loeb $ map (loebify ls m) ls
  where
    unLookup a ns = unMaybe $ lookup a ns
    ls = sort $ 0 : ((maximum nums) + 3) : nums
    m = (maximum nums + 3)

diffs :: [Int] -> [Int]
diffs [] = []
diffs [a] = []
diffs (a:b:xs) = (b-a):(diffs (b:xs))

loeb :: Functor f => f (f a -> a) -> f a
loeb x = go
  where
    go = fmap ($ go) x

loebify :: [Int] -> Int -> Int -> [(Int, Int)] -> (Int, Int)
loebify nums m t =
  if t == m
  then (\ls -> (t, 1))
  else (\ls -> (t, sum [ f ls | f <- links ]))
  where
    vals = takeWhile (<= (t + 3)) $ dropWhile (<= t) nums
    links = map lookupify vals

lookupify :: Eq a => a -> [(a,b)] -> b
lookupify k = unMaybe . (lookup k)

unMaybe :: Maybe a -> a
unMaybe Nothing = error "invalid key"
unMaybe (Just v) = v