import AoC

import Data.List

pairs :: [a] -> [(a, a)]
pairs xs = zipWith (,) xs (tail xs)

part1 :: [Int] -> Int
part1 = product . map length . group . sort . map (uncurry $ flip (-)) . pairs

numArrangements :: [Int] -> Int
numArrangements = snd . head . go
    where go xs = foldr (\x acc -> (x, sum . map snd . takeWhile (\(y, _) -> x-y <= 3) $ acc) : acc) [(last xs, 1)] $ init xs

main = runAoC (addPhonePort . sortReverse . (0:) . map read . lines) part1 numArrangements
    where addPhonePort (x:xs) = (x+3):x:xs
          sortReverse = sortBy (flip compare)