advent11/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/10/advent-of-code-2023-day-10/
   2
   3 import AoC
   4
   5 import Data.List
   6 import Data.Maybe
   7 import Linear (V2(..), (^+^), (^-^))
   8 import qualified Data.Set as S
   9
  10 type Position = V2 Int -- r, c
  11
  12 type Gals = S.Set Position
  13 data Galaxies = Galaxies Gals (Position, Position)
  14   deriving (Show)
  15
  16 main :: IO ()
  17 main =
  18   do  dataFileName <- getDataFileName
  19       text <- readFile dataFileName
  20       let galaxies = mkGalaxies text
  21       -- print galaxies
  22       -- print $ expandGalaxies galaxies 10
  23       print $ part1 galaxies
  24       print $ part2 galaxies
  25
  26 part1, part2 :: Galaxies -> Int
  27 part1 galaxies = allDistances gs
  28   where (Galaxies gs _) = expandGalaxies galaxies 2
  29
  30 part2 galaxies = allDistances gs
  31   where (Galaxies gs _) = expandGalaxies galaxies (10^6)
  32
  33 allDistances :: Gals -> Int
  34 allDistances gs = case S.minView gs of
  35   Nothing -> 0
  36   Just (g, gs') -> (S.foldl' (addDist g) 0 gs) + allDistances gs'
  37   where addDist g1 acc g2 = acc + distance g1 g2
  38
  39 distance :: Position -> Position -> Int
  40 distance g1 g2 = abs dr + abs dc
  41   where (V2 dr dc) = g1 ^-^ g2
  42
  43
  44 expandGalaxies :: Galaxies -> Int -> Galaxies
  45 expandGalaxies galaxies scale = galaxies''
  46   where er = emptyRows galaxies
  47         galaxies' = expandRows galaxies er (scale - 1)
  48         ec = emptyCols galaxies'
  49         galaxies'' = expandCols galaxies' ec (scale - 1)
  50
  51 emptyRows, emptyCols :: Galaxies -> [Int]
  52 emptyRows (Galaxies gs (V2 r0 _, V2 r1 _)) =
  53   [ r | r <- [0..r1] , S.null $ onRow gs r ]
  54   where r1 = S.findMax $ S.map (\(V2 r _) -> r) gs
  55 emptyCols (Galaxies gs (V2 _ c0, V2 _ c1)) =
  56   [ c | c <- [0..c1] , S.null $ onCol gs c ]
  57   where c1 = S.findMax $ S.map (\(V2 _ c) -> c) gs
  58
  59 onRow, onCol :: Gals -> Int -> Gals
  60 onRow gs r = S.filter (\(V2 r' _) -> r == r') gs
  61 onCol gs c = S.filter (\(V2 _ c') -> c == c') gs
  62
  63 expandRows, expandCols :: Galaxies -> [Int] -> Int -> Galaxies
  64 expandRows (Galaxies gs (V2 r0 c0, V2 r1 c1)) expansions scale =
  65   Galaxies gs' (V2 r0 c0, V2 (r1 + nExp) c1)
  66   where nExp = length expansions * scale
  67         gs' = foldr (shiftRows scale) gs expansions
  68 expandCols (Galaxies gs (V2 r0 c0, V2 r1 c1)) expansions scale =
  69   Galaxies gs' (V2 r0 c0, V2 r1 (c1 + nExp))
  70   where nExp = length expansions * scale
  71         gs' = foldr (shiftCols scale) gs expansions
  72
  73 shiftRows, shiftCols :: Int -> Int -> Gals -> Gals
  74 shiftRows scale n gs = S.union small large'
  75   where (small, large) = S.partition (\(V2 r _) -> r < n) gs
  76         large' = S.map (^+^ (V2 scale 0)) large
  77 shiftCols scale n gs = S.union small large'
  78   where (small, large) = S.partition (\(V2 _ c) -> c < n) gs
  79         large' = S.map (^+^ (V2 0 scale)) large
  80
  81 -- reading the map
  82
  83 mkGalaxies :: String -> Galaxies
  84 mkGalaxies text =
  85   Galaxies (S.fromList [ V2 r c | r <- [0..maxR], c <- [0..maxC]
  86                                       , rows !! r !! c == '#'
  87                        ])
  88            (V2 0 0, V2 maxR maxC)
  89   where rows = lines text
  90         maxR = length rows - 1
  91         maxC = (length $ head rows) - 1