b0b77bf943ee7732227544dc6e5baa82c603bf62
[advent-of-code-20.git] / 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