advent03/Main.hs

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