day 4

Signed-off-by: Amneesh Singh <natto@weirdnatto.in>

aoc2024.cabal

@@ -37,3 +37,9 @@ executable day3
   hs-source-dirs: src
   main-is:        Day3.hs
   build-depends:  libaoc
+
+executable day4
+  import:         common
+  hs-source-dirs: src
+  main-is:        Day4.hs
+  build-depends:  libaoc

lib/AoC.hs

@@ -1,10 +1,29 @@
 module AoC where
 
-import Text.Parsec
+import Data.List (transpose)
+import Text.Parsec (ParseError)
 
+-- extract Right value after parsing
 extract :: Either ParseError a -> a
 extract (Left err) = error ("Parsing failed: " ++ show err)
 extract (Right val) = val
 
+-- count elements in a list
 count :: (Eq a) => a -> [a] -> Int
 count x = length . filter (x ==)
+
+-- extract diagonals from a matrix
+diagonals :: [[a]] -> [[a]]
+diagonals =
+  (++)
+    <$> reverse . transpose . zipWith drop [0 ..]
+    <*> transpose . zipWith drop [1 ..] . transpose
+
+-- get indices of substring sub in str
+findSubstrings :: [Char] -> [Char] -> [Int]
+findSubstrings sub str = findSubstrings' sub str 0
+  where
+    findSubstrings' _ [] _ = []
+    findSubstrings' sub str@(x : xs) idx
+      | take (length sub) str == sub = idx : findSubstrings' sub xs (idx + 1)
+      | otherwise = findSubstrings' sub xs (idx + 1)

src/Day4.hs

@@ -0,0 +1,63 @@
+module Main where
+
+import qualified AoC as A (diagonals, findSubstrings)
+import Data.List (intersect, transpose)
+
+-- |
+--   A.diagonals gives the diagonals of a matrix
+--   A.findSubstrings provides the indices of a substring in a string
+--   We get the number of substrings XMAS and SAMX in all possible ways and add
+part1 :: [[Char]] -> Int
+part1 grid =
+  (sum . concat)
+    [ map countXmas grid,
+      map countXmas $ transpose grid,
+      map countXmas $ A.diagonals grid,
+      map countXmas $ A.diagonals $ map reverse grid
+    ]
+  where
+    countXmas :: [Char] -> Int
+    countXmas =
+      (+)
+        <$> length . A.findSubstrings "XMAS"
+        <*> length . A.findSubstrings "SAMX"
+
+-- |
+--   We get indices of substrings MAS and SAM across only the diagonals
+--   Then we calculate the coordinates of the letter 'A' in those
+--   Then we intersect the lists to find common As
+part2 :: [[Char]] -> Int
+part2 grid =
+  let m = length grid
+      n = length $ head grid
+      -- diagonals = left -> right, top -> bottom
+      diags1 = findMas $ A.diagonals grid
+      -- diagonals = right -> left, top -> bottom
+      diags2 = findMas $ A.diagonals $ map reverse grid
+   in length $
+        intersect
+          -- coordinate calculation for A across diag1 and diag2
+          [ if i < n
+              then (n - i + diag, diag + 1)
+              else (diag + 1, i - n + diag + 2)
+            | (i, diags) <- zip [0 ..] diags1,
+              diag <- diags
+          ]
+          [ if i < m
+              then (i - diag - 1, diag + 1)
+              else (m - diag - 2, i - m + diag + 2)
+            | (i, diags) <- zip [0 ..] diags2,
+              diag <- diags
+          ]
+  where
+    -- find indices of MAS and SAM across diagonals
+    findMas :: [[Char]] -> [[Int]]
+    findMas = map ((++) <$> A.findSubstrings "MAS" <*> A.findSubstrings "SAM")
+
+main :: IO ()
+main =
+  do
+    raw <- lines <$> readFile "./inputs/day4.in"
+
+    putStr "Part 1: " >> print (part1 raw)
+    putStr "Part 2: " >> print (part2 raw)