AoC/2021/day04/day04.hs

57 lines
1.6 KiB
Haskell
Raw Normal View History

2021-12-10 15:19:32 +01:00
import Control.Lens
import Data.List
import Data.List.Split
import Data.Maybe
type Board = [[(Int, Bool)]]
type Game = ([Board], [Int])
main :: IO ()
main = do
(i:is) <- splitOn "\n\n" <$> readFile "input"
let nums = read $ '[' : i ++ "]" :: [Int]
let boards =
map (map (flip zip (repeat False) . map read . words) . lines) is :: [[[( Int
, Bool)]]]
let (newBoards, newNums) = bingo (boards, nums)
print $ countScore newNums $ getWinBoard newBoards
let ([lb], ln) = loseBoard (newBoards, newNums)
print $ countScore ln lb
replace :: Eq a => a -> a -> [a] -> [a]
replace a b s = maybe s (f s) z
where
f x y = (element y .~ b) x
z = elemIndex a s
bingo :: Game -> Game
bingo (boards, []) = (boards, [])
bingo (boards, n:ns) =
if or $ map hasBingo next
then (next, n : ns)
else bingo (next, ns)
where
next = map (map (replace (n, False) (n, True))) boards
hasBingo :: Board -> Bool
hasBingo board = or $ map and bools ++ map and (transpose bools)
where
bools = map (map (snd)) board
getWinBoard :: [Board] -> Board
getWinBoard x = x !! (fromJust $ elemIndex True $ map hasBingo x)
countScore :: [Int] -> Board -> Int
countScore n =
(*) (head n) . sum . map fst . filter (\x -> x ^. _2 == False) . concat
loseBoard :: Game -> Game
loseBoard (boards, n) =
if null $ concat $ concat $ fst remaining
then (newBoards, newNums)
else loseBoard remaining
where
(newBoards, newNums) = bingo (boards, n)
remaining = (newBoards \\ [getWinBoard newBoards], newNums)