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