advent16/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/19/advent-of-code-2023-day-16/
   2
   3 -- import Debug.Trace
   4
   5 import AoC
   6 import Linear (V2(..), (^+^))
   7 import Data.Array.IArray
   8 import qualified Data.Set as S
   9
  10 data Element =
  11   Empty | SlashMirror | BackslashMirror |
  12   HorizontalSplitter | VerticalSplitter
  13   deriving (Show, Eq, Enum)
  14
  15 type Position = V2 Int -- r, c
  16 type Grid = Array Position Element
  17
  18 data BeamHead = BeamHead { beamPos :: Position, beamDirection :: Position }
  19   deriving (Show, Eq, Ord)
  20
  21 pattern U, D, L, R :: Position
  22 pattern U = V2 (-1) 0
  23 pattern D = V2 1 0
  24 pattern L = V2 0 (-1)
  25 pattern R = V2 0 1
  26
  27 type Energised = S.Set BeamHead
  28
  29 main :: IO ()
  30 main =
  31   do  dataFileName <- getDataFileName
  32       text <- readFile dataFileName
  33       let grid = mkGrid text
  34       -- print map
  35       print $ part1 grid
  36       print $ part2 grid
  37
  38 part1, part2 :: Grid -> Int
  39 part1 grid = countEnergised grid (BeamHead (V2 0 0) R)
  40 part2 grid = maximum $ fmap (countEnergised grid) $ getEdges grid
  41
  42 -- reading the map
  43
  44 mkGrid :: String -> Grid
  45 mkGrid 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)) $ fmap readElement $ concat rows
  50
  51 readElement :: Char -> Element
  52 readElement '/' = SlashMirror
  53 readElement '\\' = BackslashMirror
  54 readElement '|' = VerticalSplitter
  55 readElement '-' = HorizontalSplitter
  56 readElement _ = Empty
  57
  58 -- propagating the beam
  59
  60 countEnergised :: Grid -> BeamHead -> Int
  61 countEnergised grid bh = S.size $ S.map beamPos $ propagate grid S.empty [bh]
  62
  63 propagate :: Grid -> Energised -> [BeamHead] -> Energised
  64 -- propagate _ e hs | trace ("P " ++ (show hs)) False = undefined
  65 propagate _ energised [] = energised
  66 propagate grid energised (bh:bhs)
  67   | S.member bh energised = propagate grid energised bhs
  68   | otherwise = propagate grid energised' (bhs ++ nexts')
  69   where this = grid ! (beamPos bh)
  70         nexts = propagateElem this bh
  71         nexts' = filter ((inRange (bounds grid)) . beamPos) nexts
  72         energised' = S.insert bh energised
  73
  74 propagateElem :: Element -> BeamHead -> [BeamHead]
  75 -- propagateElem e h | trace ("PE " ++ (show e) ++ " " ++ (show h)) False = undefined
  76 propagateElem Empty (BeamHead pos dir) = [BeamHead (pos ^+^ dir) dir]
  77 propagateElem SlashMirror (BeamHead pos L) = [BeamHead (pos ^+^ D) D]
  78 propagateElem SlashMirror (BeamHead pos R) = [BeamHead (pos ^+^ U) U]
  79 propagateElem SlashMirror (BeamHead pos U) = [BeamHead (pos ^+^ R) R]
  80 propagateElem SlashMirror (BeamHead pos D) = [BeamHead (pos ^+^ L) L]
  81 propagateElem BackslashMirror (BeamHead pos L) = [BeamHead (pos ^+^ U) U]
  82 propagateElem BackslashMirror (BeamHead pos R) = [BeamHead (pos ^+^ D) D]
  83 propagateElem BackslashMirror (BeamHead pos U) = [BeamHead (pos ^+^ L) L]
  84 propagateElem BackslashMirror (BeamHead pos D) = [BeamHead (pos ^+^ R) R]
  85 propagateElem HorizontalSplitter (BeamHead pos L) = [BeamHead (pos ^+^ L) L]
  86 propagateElem HorizontalSplitter (BeamHead pos R) = [BeamHead (pos ^+^ R) R]
  87 propagateElem HorizontalSplitter (BeamHead pos _) =
  88   [BeamHead (pos ^+^ L) L, BeamHead (pos ^+^ R) R]
  89 propagateElem VerticalSplitter (BeamHead pos U) = [BeamHead (pos ^+^ U) U]
  90 propagateElem VerticalSplitter (BeamHead pos D) = [BeamHead (pos ^+^ D) D]
  91 propagateElem VerticalSplitter (BeamHead pos _) =
  92   [BeamHead (pos ^+^ U) U, BeamHead (pos ^+^ D) D]
  93
  94 getEdges :: Grid -> [BeamHead]
  95 getEdges grid = [BeamHead (V2 0 c) D | c <- [0..maxC]] ++
  96                 [BeamHead (V2 r 0) R | r <- [0..maxR]] ++
  97                 [BeamHead (V2 r maxC) L | r <- [0..maxR]] ++
  98                 [BeamHead (V2 maxR c) U | c <- [0..maxC]]
  99   where (V2 maxR maxC) = snd $ bounds grid