advent23/Main.hs

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