2024 Day Four

DayFour.hs

{-# LANGUAGE OverloadedStrings #-}
module DayFour (allPaths, wordSearch, listTo2dArray, sampleXmas, addLookahead, partTwo) where

import qualified Data.Text as T
import Data.Array
import Data.Set (fromList, intersection, Set)
import qualified Data.List.ZigZag as Z
import Data.List (transpose)

sampleXmas :: [String]
sampleXmas = [
    ".M.S......",
    "..A..MSMS.",
    ".M.S.MAA..",
    "..A.ASMSM.",
    ".M.S.M....",
    "..........",
    "S.S.S.S.S.",
    ".A.A.A.A..",
    "M.M.M.M.M.",
    ".........."]

allPaths :: [T.Text] -> [T.Text]
allPaths rows = rows ++ columns ++ diagonals ++ reverseDiagonals
  where unpacked = map T.unpack rows
        reverseUnpacked = map (T.unpack . T.reverse) rows
        columnsUnpacked = transpose unpacked
        columns = map T.pack columnsUnpacked
        diagonals = map T.pack $ Z.diagonals unpacked
        reverseDiagonals = map T.pack $ Z.diagonals reverseUnpacked

wordSearch :: T.Text -> [T.Text] -> Int
wordSearch target rows = sum (map (T.count target) (allPaths rows)) + sum (map (T.count (T.reverse target)) (allPaths rows))



-- Part 2:
mkArray :: (Ix a) => (a -> b) -> (a,a) -> Array a b
mkArray f bnds =  array bnds [(i, f i) | i <- range bnds]

addLookahead :: [[Char]] -> [[(Char, Char)]]
addLookahead = map withLookAhead 
    where withLookAhead xs = zip xs $ drop 2 xs

listTo2dArray :: [[t]] -> Array (Int, Int) t
listTo2dArray grid = mkArray findInGrid ((0,0), (length grid - 1, length (head grid) - 1)) 
  where findInGrid (row, col) = (grid !! row) !! col

examineOneOrientation :: [String] -> Set (Int, Int)
examineOneOrientation grid = indicesOfAs `intersection` doubleMpoints `intersection` doubleSpoints
  where indicesOfAs = fromList $ [idx | (idx, c) <- assocs asArray, c == 'A']
        asArray = listTo2dArray grid
        withLookahead = listTo2dArray $ addLookahead grid
        doubleMpoints = fromList $ concatMap (yieldPoints ('M','M')) $ assocs withLookahead
        doubleSpoints = fromList $ concatMap (yieldPoints ('S','S')) $ assocs withLookahead
        yieldPoints target ((r,c), actual) 
          | target == actual = [(r+1, c+1), (r-1, c+1)]
          | otherwise = []

partTwo :: [String] -> Int
partTwo grid = length (examineOneOrientation grid ) + length (examineOneOrientation ( transpose grid))