advent03/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/14/advent-of-code-2022-day-12/
   2
   3 import AoC
   4
   5
   6 import Data.Char
   7 import Data.List
   8 import Linear (V2(..), (^+^))
   9 import Data.Array.IArray
  10
  11 type Position = V2 Int -- r, c
  12 type Engine = Array Position Char
  13 type NumPos = [Position]
  14
  15 data NumberSeek =
  16   NumberSeek { inNumber :: Bool
  17              , positions :: NumPos
  18              , foundNumbers :: [NumPos]
  19              } deriving (Show)
  20
  21
  22 main :: IO ()
  23 main =
  24   do  dataFileName <- getDataFileName
  25       text <- readFile dataFileName
  26       let engine = mkEngine text
  27       let allNums = findNumbers engine
  28       let partNums = filter (numTouchingSymbols engine) allNums
  29       -- print engine
  30       -- print $ findNumbers engine
  31       print $ part1 engine partNums
  32       print $ part2 engine partNums
  33
  34
  35 part1, part2 :: Engine -> [NumPos] -> Int
  36 part1 engine partNums = sum partNumValues
  37   where partNumValues = map (readNumber engine) partNums
  38
  39 part2 engine partNums = sum $ fmap product gearRatios
  40   where sts = stars engine
  41         touchingStars = fmap (adjacentNumbers partNums) sts
  42         gears = filter ((==) 2 . length) touchingStars
  43         gearRatios = [[readNumber engine n | n <- pNs] | pNs <- gears]
  44
  45 mkEngine :: String -> Engine
  46 mkEngine text = grid
  47   where rows = lines text
  48         r = length rows - 1
  49         c = (length $ head rows) - 1
  50         grid = listArray ((V2 0 0), (V2 r c)) $ concat rows
  51
  52 isEngineSymbol :: Char -> Bool
  53 isEngineSymbol c = (not $ isDigit c) && (c /= '.')
  54
  55 findNumbers :: Engine -> [NumPos]
  56 findNumbers engine = numbers
  57   where ((V2 r1 _), (V2 r2 _)) = bounds engine
  58         rows = [r1..r2]
  59         numbers = concatMap (foundNumbers . (findNumbersInRow engine)) rows
  60
  61 findNumbersInRow :: Engine -> Int -> NumberSeek
  62 findNumbersInRow engine r
  63   | inNumber finalSeek = NumberSeek False [] ((reverse $ positions finalSeek):foundNumbers finalSeek)
  64   | otherwise = finalSeek
  65   where finalSeek = foldl' (buildNumber engine)
  66                            (NumberSeek False [] [])
  67                            $ range $ rowBounds engine r
  68
  69 buildNumber :: Engine -> NumberSeek -> Position -> NumberSeek
  70 buildNumber engine NumberSeek{..} p
  71   | inNumber && isDigit c = NumberSeek True (p:positions) foundNumbers
  72   | inNumber && not (isDigit c) = NumberSeek False [] ((reverse positions):foundNumbers)
  73   | not inNumber && isDigit c = NumberSeek True [p] foundNumbers
  74   | otherwise = NumberSeek False [] foundNumbers
  75   where c = engine ! p
  76
  77 rowBounds :: Engine -> Int -> (V2 Int, V2 Int)
  78 rowBounds engine r = (V2 r c1, V2 r c2)
  79   where ((V2 _ c1), (V2 _ c2)) = bounds engine
  80
  81 neighbours :: Position -> [Position]
  82 neighbours p = [p ^+^ V2 dr dc | dr <- [-1..1], dc <- [-1..1]
  83                                , (dr, dc) /= (0, 0) ]
  84
  85 touchingSymbol :: Engine -> Position -> Bool
  86 touchingSymbol engine p = any (isEngineSymbol . (engine !)) nbrs
  87   where nbrs = filter (inRange (bounds engine)) $ neighbours p
  88
  89 numTouchingSymbols :: Engine -> NumPos -> Bool
  90 numTouchingSymbols engine ps = any (touchingSymbol engine) ps
  91
  92 readNumber :: Engine -> NumPos -> Int
  93 readNumber engine ps = read $ map (engine !) ps
  94
  95 stars :: Engine -> [Position]
  96 stars engine = filter ((==) '*' . (engine !)) $ indices engine
  97
  98 starTouchesNumber :: Position -> NumPos -> Bool
  99 starTouchesNumber star ps = not $ null $ intersect ps halo
 100   where halo = neighbours star
 101
 102 adjacentNumbers :: [NumPos] -> Position -> [NumPos]
 103 adjacentNumbers nums star = filter (starTouchesNumber star) nums