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       print $ part1 galaxies
  22       print $ part2 galaxies
  23
  24 part1, part2 :: Galaxies -> Int
  25 part1 = allDistances . expandGalaxies 2
  26 part2 = allDistances . expandGalaxies (10^6)
  27
  28 allDistances :: Galaxies -> Int
  29 allDistances gs = case S.minView gs of
  30   Nothing -> 0
  31   Just (g, gs') -> (S.foldl' (addDist g) 0 gs) + allDistances gs'
  32   where addDist g1 acc g2 = acc + distance g1 g2
  33
  34 distance :: Position -> Position -> Int
  35 distance g1 g2 = abs dr + abs dc
  36   where (V2 dr dc) = g1 ^-^ g2
  37
  38
  39 expandGalaxies :: Int -> Galaxies -> Galaxies
  40 expandGalaxies scale galaxies = galaxies''
  41   where er = emptyRows galaxies
  42         galaxies' = expandRows galaxies er scale
  43         ec = emptyCols galaxies'
  44         galaxies'' = expandCols galaxies' ec scale
  45
  46 emptyRows, emptyCols :: Galaxies -> [Int]
  47 emptyRows galaxies = [ r | r <- [0..r1] , S.null $ onRow galaxies r ]
  48   where r1 = S.findMax $ S.map (\(V2 r _) -> r) galaxies
  49 emptyCols galaxies = [ c | c <- [0..c1] , S.null $ onCol galaxies c ]
  50   where c1 = S.findMax $ S.map (\(V2 _ c) -> c) galaxies
  51
  52 onRow, onCol :: Galaxies -> Int -> Galaxies
  53 onRow galaxies r = S.filter (\(V2 r' _) -> r == r') galaxies
  54 onCol galaxies c = S.filter (\(V2 _ c') -> c == c') galaxies
  55
  56 expandRows, expandCols :: Galaxies -> [Int] -> Int -> Galaxies
  57 expandRows galaxies expansions scale = foldr (shiftRows scale) galaxies expansions
  58 expandCols galaxies expansions scale = foldr (shiftCols scale) galaxies expansions
  59
  60 shiftRows, shiftCols :: Int -> Int -> Galaxies -> Galaxies
  61 shiftRows scale n galaxies = S.union small large'
  62   where (small, large) = S.partition (\(V2 r _) -> r < n) galaxies
  63         large' = S.map (^+^ (V2 (scale - 1) 0)) large
  64 shiftCols scale n galaxies = S.union small large'
  65   where (small, large) = S.partition (\(V2 _ c) -> c < n) galaxies
  66         large' = S.map (^+^ (V2 0 (scale - 1))) large
  67
  68 -- reading the map
  69
  70 mkGalaxies :: String -> Galaxies
  71 mkGalaxies text = S.fromList [ V2 r c | r <- [0..maxR], c <- [0..maxC]
  72                                       , rows !! r !! c == '#'
  73                              ]
  74   where rows = lines text
  75         maxR = length rows - 1
  76         maxC = (length $ head rows) - 1