adventofcode-2020/day9/main.hs

#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple containers])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports #-}
{-# LANGUAGE OverloadedStrings #-}

import Data.IntMap.Strict (IntMap)
import Data.IntSet (IntSet)
import Data.List (inits, tails, sortOn)
import Data.Monoid
import Debug.Trace (trace)
import Text.Pretty.Simple
import qualified Data.IntMap.Strict as M
import qualified Data.IntSet as S


exampleData :: [Int]
exampleData = [ 35,20,15,25,47,40,62,55,65,95,102,117,150
              , 182,127,219,299,277,309,576 ];

pairSums :: [Int] -> IntSet
pairSums xs = S.fromList $ do
  i1 <- xs
  i2 <- xs
  pure (i1 + i2)

isSumOfPreamble :: Int -> [Int] -> Bool
isSumOfPreamble size xs = go tx
  where
    go (x:_)  = x `S.member` (pairSums preamble)
    go ([])   = False
    (preamble, tx) = splitAt size xs

solvePart1 :: Int -> [Int] -> (Int,Bool)
solvePart1 size message =
  head $ dropWhile (snd) $
  map go ( takeWhile ((> size) . length) $ tails message )
  where
    go xs = (head $ drop size xs,isSumOfPreamble size xs)


--------------------------------------------------------------------------------

sortedContiguousSubLists :: [Int] -> [[Int]]
sortedContiguousSubLists =
  dropWhile (== []) . sortOn length . concat . fmap inits . tails

listSums :: [[Int]] -> IntMap [Int]
listSums xss = M.fromList $ do
  xs <- xss
  pure $ (getSum $ foldMap Sum xs, xs)

solvePart2 :: Int -> [Int] -> Int
solvePart2 target message = (\x -> minimum x + maximum x) $  bigMap M.! target
  where
    bigMap = listSums $ sortedContiguousSubLists message

main :: IO ()
main = do
  putStrLn "Test"
  print exampleData
  print $ 1 `S.member` (pairSums [])
  print $ pairSums exampleData
  putStrLn "Day 9 - Part 1"
  print $ solvePart1 5 exampleData
  input <- lines <$> readFile "day9/input"
  print $ solvePart1 25 $ map read input
  putStrLn "Test"
  print $ take 10 $ drop 20 $ sortedContiguousSubLists exampleData
  putStrLn "Day 9 - Part 2"
  print $ solvePart2 127 exampleData
  print $ solvePart2 3199139634 $ map read input