advent10/src/advent10.hs

   1 -- import Debug.Trace
   2
   3 import qualified Data.Map as M
   4 -- import qualified Data.Set as S
   5 import Data.List
   6
   7 main :: IO ()
   8 main =
   9     do
  10         numStrs <- readFile "data/advent10.txt"
  11         let nums = map (read @Int) $ lines numStrs
  12         let device = 3 + maximum nums
  13         let sortedNums = sort (0:device:nums)
  14         print $ part1 sortedNums
  15         print $ part2 sortedNums
  16
  17 part1 :: [Int] -> Int
  18 part1 nums = jolt1s * jolt3s
  19     where jolt1s = joltNs 1 nums
  20           jolt3s = joltNs 3 nums
  21
  22 joltNs d ns = length $ filter diff $ zip ns (tail ns)
  23     where diff (a, b) = (b - a) == d
  24
  25 part2 :: [Int] -> Int
  26 part2 nums = snd $ M.findMax $ foldl includeAdapter (M.singleton 0 1) $ tail nums
  27
  28 includeAdapter arrangements adapter = M.insert adapter earlierCount arrangements
  29     where
  30         earlierArrangements = M.filterWithKey (\k _ -> (adapter - k) <= 3) arrangements
  31         earlierCount = M.foldl (+) 0 earlierArrangements