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