AoC/2021/day04/day04.hs

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)