advent-2021/4/solution.hs

{-# LANGUAGE TupleSections #-}
import Data.List
import Debug.Trace

main :: IO ()
main = readFile "4/input.txt" >>= print . solve . parse . lines

type Input = ([Int], [Board])

parse :: [String] -> Input
parse (l : ls) = (n, b)
    where n = map read $ words [ if c == ',' then ' ' else c | c <- l ]
          b = parseBoards (filter (/= "") ls)

parseBoards :: [String] -> [Board]
parseBoards [] = []
parseBoards l  = parse a : parseBoards b
    where (a,b) = splitAt 5 l 
          parse = map $ map (Unmarked . read) . words

type Board = [[Cell]]

rows :: Board -> [[Cell]]
rows = id

cols :: Board -> [[Cell]]
cols = transpose

hasWon :: Board -> Bool
hasWon b = check (rows b) || check (cols b)
    where check = any $ all marked

score :: Int -> Board -> Int
score n b = n * sum [ x | c <- b, Unmarked x <- c ]

call :: Int -> Board -> Board
call n = map (map mark)
    where mark (Unmarked x) | x == n = Marked x
          mark x = x

data Cell = Unmarked Int
          | Marked Int

marked :: Cell -> Bool
marked (Marked _) = True
marked _          = False

solve :: Input -> (Int, Int)
solve input = (f input, g input)
    where wins extr (n,b) = uncurry score $ extr $ bingo n b
          f = wins head
          g = wins last

bingo :: [Int] -> [Board] -> [(Int, Board)]
bingo _ []              = []
bingo (n : ns) (b : bs) = map (n,) w ++ bingo ns l
    where (w,l) = partition hasWon r
          r     = iter n (b : bs)

iter :: Int -> [Board] -> [Board]
iter _ []       = []
iter n (b : bs) = call n b : iter n bs