advent23/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/16/advent-of-code-2021-day-15/
   2
   3 import Debug.Trace
   4
   5
   6 import qualified Data.PQueue.Prio.Min as P
   7 import qualified Data.Set as S
   8 import qualified Data.Sequence as Q
   9 import Data.Sequence ((<|), (|>), (><)) --, ViewR( (:>) ), ViewL( (:<) ))
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import Data.Foldable (foldl', sum) -- (toList, foldr', foldl', all)
  13 -- import Data.Char
  14 import Control.Monad.Reader
  15 import Control.Lens hiding ((<|), (|>), (:>), (:<))
  16 import Data.Maybe -- (fromMaybe)
  17 -- import Linear (V2(..), (^+^)) --, (^-^), (*^), (^*))
  18 import Linear hiding (trace)
  19
  20
  21
  22 pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
  23 pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|)
  24 pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>)
  25
  26
  27 data Amphipod = A | B | C | D deriving (Show, Read, Eq, Ord, Enum)
  28
  29 type Coord = V2 Int -- r, c
  30 _r :: Lens' (V2 Int) Int
  31 _r = _x
  32 _c :: Lens' (V2 Int) Int
  33 _c = _y
  34
  35 data Step = Step
  36   { _destination :: Coord
  37   , _distance :: Int
  38   , _transits :: S.Set Coord
  39   , _entryRequirement :: Maybe Amphipod
  40   } deriving (Show, Eq, Ord)
  41 makeLenses ''Step
  42
  43 type Steps = M.Map Coord (S.Set Step)
  44
  45 data Burrow = Burrow
  46   { _possibleSteps :: Steps
  47   , _roomColumns :: M.Map Amphipod Int
  48   , _hallRow :: Int
  49   } deriving (Show, Eq)
  50 makeLenses ''Burrow
  51
  52 type BurrowContext = Reader Burrow
  53
  54 type MoveState = M.Map Coord Amphipod
  55
  56 data AppliedMove = AppliedMove
  57   { _afterMove :: MoveState
  58   , _appliedStep :: Step
  59   }
  60   deriving (Show, Eq, Ord)
  61 makeLenses ''AppliedMove
  62
  63 data Agendum =
  64     Agendum { _current :: MoveState
  65             , _trail :: Q.Seq MoveState
  66             , _trailCost :: Int
  67             , _cost :: Int
  68             } deriving (Show, Eq)
  69 makeLenses ''Agendum
  70
  71 type Agenda = P.MinPQueue Int Agendum
  72
  73 type ExploredStates = S.Set MoveState
  74
  75
  76
  77 main :: IO ()
  78 main =
  79   do  text <- readFile "data/advent23.txt"
  80       -- let (burrow, initState) = mkBurrow text
  81       -- print burrow
  82       -- print initState
  83       print $ part1 text
  84       print $ part2 text
  85
  86
  87 -- part1 :: Burrow -> MoveState -> Int
  88 part1 text = maybe 0 _cost result
  89     where
  90       (burrow, initState) = mkBurrow text
  91       result = runReader (searchBurrow initState) burrow
  92
  93 part2 text = maybe 0 _cost result
  94     where
  95       rows = lines text
  96       extraRows = [("  #D#C#B#A#  " :: String), ("  #D#B#A#C#  " :: String)]
  97       modifiedRows = (take 3 rows) ++ extraRows ++ (drop 3 rows)
  98       modifiedText = unlines modifiedRows
  99       (burrow, initState) = mkBurrow modifiedText
 100       result = runReader (searchBurrow initState) burrow
 101
 102
 103 searchBurrow :: MoveState -> BurrowContext (Maybe Agendum)
 104 searchBurrow initState =
 105     do agenda <- initAgenda initState
 106        aStar agenda S.empty
 107
 108 initAgenda :: MoveState -> BurrowContext Agenda
 109 initAgenda initState =
 110     do c <- estimateCost initState
 111        return $ P.singleton c Agendum { _current = initState
 112                                       , _trail = Q.empty, _trailCost = 0, _cost = c}
 113
 114
 115 aStar ::  Agenda -> ExploredStates -> BurrowContext (Maybe Agendum)
 116 aStar agenda closed
 117     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 118     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 119     | P.null agenda = return Nothing
 120     | otherwise =
 121         do  let (_, currentAgendum) = P.findMin agenda
 122             let reached = currentAgendum ^. current
 123             nexts <- candidates currentAgendum closed
 124             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 125             reachedGoal <- isGoal reached
 126             if reachedGoal
 127             then return (Just currentAgendum)
 128             else if reached `S.member` closed
 129                  then aStar (P.deleteMin agenda) closed
 130                  else aStar newAgenda (S.insert reached closed)
 131
 132
 133 candidates ::  Agendum -> ExploredStates -> BurrowContext (Q.Seq Agendum)
 134 candidates agendum closed =
 135     do  let candidate = agendum ^. current
 136         let previous = agendum ^. trail
 137         let prevCost = agendum ^. trailCost
 138         succs <- successors candidate
 139         let nonloops = S.filter (\s -> (s ^. afterMove) `S.notMember` closed) succs
 140         let nonloopsQ = Q.fromList $ S.toList nonloops
 141         mapM (makeAgendum previous prevCost) nonloopsQ
 142
 143
 144 makeAgendum ::  Q.Seq MoveState -> Int -> AppliedMove -> BurrowContext Agendum
 145 makeAgendum previous prevCost newPosition =
 146     do predicted <- estimateCost (newPosition ^. afterMove)
 147        -- grid <- asks _grid
 148        let newTrail = previous |> (newPosition ^. afterMove)
 149        let newPositionCost = stepCost newPosition
 150        let incurred = prevCost + newPositionCost
 151        return Agendum { _current = newPosition ^. afterMove
 152                       , _trail = newTrail
 153                       , _trailCost = incurred
 154                       , _cost = incurred + predicted
 155                       }
 156
 157 -- class (Eq s, Ord s, Show s) => SearchState s where
 158 --     emptySearchState :: MoveState
 159 --     successors :: MoveState -> BurrowContext (Q.Seq MoveState)
 160 --     estimateCost :: MoveState -> BurrowContext Int
 161 --     isGoal :: MoveState -> BurrowContext Bool
 162 --     entryCost :: MoveState -> BurrowContext Int
 163
 164
 165 -- instance SearchState Position where
 166
 167 --   emptySearchState = Position (V2 0 0)
 168
 169 successors :: MoveState -> BurrowContext (S.Set AppliedMove)
 170 successors moveState =
 171   do steps <- asks _possibleSteps
 172      let succs = M.foldrWithKey' (legalSteps steps moveState) S.empty moveState
 173      return succs
 174
 175 legalSteps :: Steps -> MoveState -> Coord -> Amphipod -> S.Set AppliedMove -> S.Set AppliedMove
 176 legalSteps steps state here amphipod acc = S.union appliedSteps acc
 177   where allSteps = steps ! here
 178         freeSteps = S.filter freeSpaces allSteps
 179         freeSpaces st = S.null $ S.intersection (M.keysSet state) (st ^. transits)
 180         validTargetSteps = S.filter (\st -> fromMaybe amphipod (st ^. entryRequirement) == amphipod) freeSteps
 181         openRoomSteps = S.filter (openRoom state) validTargetSteps
 182         appliedSteps = S.map (\s -> AppliedMove
 183                                       { _afterMove = (applyStep state here s)
 184                                       , _appliedStep =  s
 185                                       }
 186                               ) openRoomSteps
 187
 188 openRoom :: MoveState -> Step -> Bool
 189 openRoom state step
 190   | isNothing e = True
 191   | otherwise = M.null roomBlockers
 192   where e = step ^. entryRequirement
 193         je = fromJust e
 194         tc = step ^. destination . _c
 195         roomBlockers = M.filterWithKey (\(V2 _ ac) a -> a /= je && ac == tc) state
 196
 197 applyStep :: MoveState -> Coord -> Step -> MoveState
 198 applyStep moveState here step = moveState''
 199   where moveState' = M.delete here moveState
 200         moveState'' = M.insert (step ^. destination) (moveState ! here) moveState'
 201
 202 singleStepCost :: Amphipod -> Int
 203 singleStepCost A = 1
 204 singleStepCost B = 10
 205 singleStepCost C = 100
 206 singleStepCost D = 1000
 207
 208 estimateCost :: MoveState -> BurrowContext Int
 209 estimateCost state =
 210   do rCols <- asks _roomColumns
 211      hRow <- asks _hallRow
 212      let amphipodCosts = M.mapWithKey (estimateACost rCols hRow) state
 213      return $ sum $ M.elems amphipodCosts
 214
 215 estimateACost :: M.Map Amphipod Int -> Int -> Coord -> Amphipod -> Int
 216 estimateACost rCols hRow (V2 r c) amphipod = (singleStepCost amphipod) * dist
 217     where targetCol = rCols ! amphipod
 218           dist = if c == targetCol
 219                  then 0
 220                  else (r - hRow) + (abs (c - targetCol)) + 1
 221
 222 stepCost :: AppliedMove -> Int
 223 stepCost aStep = (singleStepCost amphipod) * (S.size $ aStep ^. appliedStep . transits)
 224   where dest = aStep ^. appliedStep . destination
 225         amphipod = (aStep ^. afterMove) ! dest
 226
 227 isGoal :: MoveState -> BurrowContext Bool
 228 isGoal state =
 229   do rCols <- asks _roomColumns
 230      let misplaced = M.filterWithKey (inWrongRoom rCols) state
 231      return $ M.null misplaced
 232
 233 inWrongRoom rCols (V2 _ c) amphipod = c /= rightCol
 234   where rightCol = rCols ! amphipod
 235
 236
 237 ------------------------------
 238
 239
 240 mkBurrow :: String -> (Burrow, MoveState)
 241 -- mkBurrow :: String -> ((S.Set Coord, M.Map Coord Amphipod), MoveState)
 242 mkBurrow text = (burrow, initState) -- (burrow, initState)
 243   where rows = lines text
 244         hall = mkHall (rows!!1)
 245         rooms = mkRooms $ drop 2 rows
 246         roomCols = S.map (^. _c) $ M.keysSet rooms
 247         hall' = S.filter ((`S.notMember` roomCols) . (^. _c)) hall
 248         routes = mkRoutes hall' rooms
 249         roomColMap = M.fromList $ zip [A .. D] $ S.toAscList roomCols
 250         burrow = Burrow { _possibleSteps = routes, _roomColumns = roomColMap, _hallRow = 1}
 251         initState = mkInitialState rows
 252
 253
 254 mkHall :: String -> S.Set Coord
 255 mkHall text = S.fromList hallCoords
 256   where hallCols = filter ((/= '#') . snd) $ zip [0..] text
 257         hallCoords = map ((V2 1) . fst) hallCols
 258
 259 mkRooms :: [String] -> M.Map Coord Amphipod
 260 mkRooms text = M.unions rooms
 261   where rooms = map mkRoom $ zip [2..] text
 262
 263 mkRoom :: (Int, String) -> M.Map Coord Amphipod
 264 mkRoom (r, text) = M.fromList roomCoords
 265   where roomCols = filter ((`elem` ("ABCD." :: String)) . snd) $ zip [0..] text
 266         roomCoords = zip (map ((V2 r) . fst) roomCols) [A .. D]
 267
 268 -- invertRooms rooms = M.fromList [(a, M.keysSet $ M.filter (== a) rooms) | a <- [A .. D]]
 269
 270 mkRoutes :: S.Set Coord -> M.Map Coord Amphipod -> Steps
 271 mkRoutes halls rooms = M.unionsWith (S.union) [hallRoutes, roomHallRoutes, roomRoomRoutes]
 272   where hallRoutes = S.foldr' (mkHallRoute rooms) M.empty halls
 273         roomHallRoutes = S.foldr' (mkRoomHallRoute halls) M.empty (M.keysSet rooms)
 274         roomRoomRoutes = S.foldr' (mkRoomRoomRoute hallRow rooms) M.empty (M.keysSet rooms)
 275         hallRow = (S.findMin halls) ^. _r
 276
 277 mkHallRoute :: M.Map Coord Amphipod -> Coord -> Steps -> Steps
 278 -- mkHallRoute rooms here routes | trace ("mkHR " ++ (show here) ++ " " ++ (show routes)) False = undefined
 279 mkHallRoute rooms here routes = M.foldrWithKey' (mkHallRoute1 here) routes rooms
 280
 281 mkHallRoute1 :: Coord -> Coord -> Amphipod -> Steps -> Steps
 282 -- mkHallRoute1 here there entry routes | trace ("mkHR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes)) False = undefined
 283 mkHallRoute1 here@(V2 hr hc) there@(V2 tr tc) entry routes = M.insert here (S.insert step existingRoutes) routes
 284   -- | trace ("mkHR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes) ++ " > " ++ show res) False = undefined
 285   -- | otherwise = res
 286   where step = Step { _destination = there
 287                     , _distance = (S.size transits)
 288                     , _transits = transits
 289                     , _entryRequirement = Just entry
 290                     }
 291         cMin = min hc tc
 292         cMax = max hc tc
 293         transits = S.delete here $ S.fromList $ [V2 hr c | c <- [cMin..cMax]] ++ [V2 r tc | r <- [hr..tr]]
 294         existingRoutes = M.findWithDefault S.empty here routes
 295         -- res = M.insert here (S.insert step existingRoutes) routes
 296
 297 mkRoomHallRoute :: S.Set Coord -> Coord -> Steps -> Steps
 298 mkRoomHallRoute halls here routes = S.foldr' (mkRoomHallRoute1 here) routes halls
 299
 300 mkRoomHallRoute1 :: Coord -> Coord -> Steps -> Steps
 301 mkRoomHallRoute1 here@(V2 hr hc) there@(V2 tr tc) routes = M.insert here (S.insert step existingRoutes) routes
 302     where step = Step { _destination = there
 303                       , _distance = (S.size transits)
 304                       , _transits = transits
 305                       , _entryRequirement = Nothing
 306                       }
 307           cMin = min hc tc
 308           cMax = max hc tc
 309           transits = S.delete here $ S.fromList $ [V2 r hc | r <- [tr..hr]] ++ [V2 tr c | c <- [cMin..cMax]]
 310           existingRoutes = M.findWithDefault S.empty here routes
 311
 312 mkRoomRoomRoute :: Int -> M.Map Coord Amphipod -> Coord -> Steps -> Steps
 313 mkRoomRoomRoute hallRow rooms here routes =  M.foldrWithKey' (mkRoomRoomRoute1 hallRow here) routes rooms
 314
 315 mkRoomRoomRoute1 :: Int -> Coord -> Coord -> Amphipod -> Steps -> Steps
 316 -- mkRoomRoomRoute1 _hallRow here there entry routes | trace ("mkRR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes)) False = undefined
 317 mkRoomRoomRoute1 hallRow here@(V2 hr hc) there@(V2 tr tc) entry routes
 318   | here == there = routes
 319   | otherwise = M.insert here (S.insert step existingRoutes) routes
 320   where step = Step { _destination = there
 321                     , _distance = (S.size transits)
 322                     , _transits = transits
 323                     , _entryRequirement = Just entry
 324                     }
 325         cMin = min hc tc
 326         cMax = max hc tc
 327         transitUp = S.fromList [V2 r hc | r <- [hallRow..hr]]
 328         transitAcross = S.fromList [V2 hallRow c | c <- [cMin..cMax]]
 329         transitDown = S.fromList [V2 r tc | r <- [hallRow..tr]]
 330         transits = S.delete here $ S.unions [transitUp, transitAcross, transitDown]
 331         existingRoutes = M.findWithDefault S.empty here routes
 332
 333
 334 mkInitialState :: [String] -> MoveState
 335 mkInitialState rows =
 336   M.fromList [ (V2 r c, read [(rows!!r)!!c])
 337              | r <- [0..maxR], c <- [0..maxC]
 338              , isAmphipod ((rows!!r)!!c)
 339              ]
 340   where maxR = length rows - 1
 341         maxC = (length $ head rows) - 1
 342         isAmphipod c = c `elem` ("ABCD" :: String)
 343