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 :: [Board] 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)