advent11/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/11/advent-of-code-2023-day-11/
   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 Galaxies = S.Set Position
  13
  14 main :: IO ()
  15 main =
  16   do  dataFileName <- getDataFileName
  17       text <- readFile dataFileName
  18       let galaxies = mkGalaxies text
  19       -- print galaxies
  20       -- print $ expandGalaxies galaxies 10
  21
  22       -- print $ allDist $ expandGalaxies 2 galaxies
  23       -- print $ allDist $ expandGalaxies (10^6) galaxies
  24
  25       print $ part1 galaxies
  26       print $ part2 galaxies
  27
  28 part1, part2 :: Galaxies -> Int
  29 part1 = allDistances . expandGalaxies 2
  30 part2 = allDistances . expandGalaxies 10e6
  31
  32 allDistances :: Galaxies -> Int
  33 allDistances gs = snd $ S.foldl' addGalaxy (S.empty, 0) gs
  34   where addGalaxy (seen, d) new =
  35             (S.insert new seen, S.foldl' (addDist new) d seen)
  36         addDist g1 d g2 = d + distance g1 g2
  37
  38 distance :: Position -> Position -> Int
  39 distance g1 g2 = abs dr + abs dc
  40   where (V2 dr dc) = g1 ^-^ g2
  41
  42
  43 expandGalaxies :: Int -> Galaxies -> Galaxies
  44 expandGalaxies scale galaxies = galaxies''
  45   where er = emptyRows galaxies
  46         galaxies' = expandRows galaxies er scale
  47         ec = emptyCols galaxies'
  48         galaxies'' = expandCols galaxies' ec scale
  49
  50 emptyRows, emptyCols :: Galaxies -> [Int]
  51 emptyRows galaxies = [ r | r <- [0..r1] , S.null $ onRow galaxies r ]
  52   where r1 = S.findMax $ S.map (\(V2 r _) -> r) galaxies
  53 emptyCols galaxies = [ c | c <- [0..c1] , S.null $ onCol galaxies c ]
  54   where c1 = S.findMax $ S.map (\(V2 _ c) -> c) galaxies
  55
  56 onRow, onCol :: Galaxies -> Int -> Galaxies
  57 onRow galaxies r = S.filter (\(V2 r' _) -> r == r') galaxies
  58 onCol galaxies c = S.filter (\(V2 _ c') -> c == c') galaxies
  59
  60 expandRows, expandCols :: Galaxies -> [Int] -> Int -> Galaxies
  61 expandRows galaxies expansions scale = foldr (shiftRows scale) galaxies expansions
  62 expandCols galaxies expansions scale = foldr (shiftCols scale) galaxies expansions
  63
  64 shiftRows, shiftCols :: Int -> Int -> Galaxies -> Galaxies
  65 shiftRows scale n galaxies = S.union small large'
  66   where (small, large) = S.partition (\(V2 r _) -> r < n) galaxies
  67         large' = S.map (^+^ (V2 (scale - 1) 0)) large
  68 shiftCols scale n galaxies = S.union small large'
  69   where (small, large) = S.partition (\(V2 _ c) -> c < n) galaxies
  70         large' = S.map (^+^ (V2 0 (scale - 1))) large
  71
  72 -- reading the map
  73
  74 mkGalaxies :: String -> Galaxies
  75 mkGalaxies text = S.fromList [ V2 r c | r <- [0..maxR], c <- [0..maxC]
  76                                       , rows !! r !! c == '#'
  77                              ]
  78   where rows = lines text
  79         maxR = length rows - 1
  80         maxC = (length $ head rows) - 1