1 -- Writeup at https://work.njae.me.uk/2022/01/03/advent-of-code-2021-day-23/
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 ((<|), (|>), (:>), (:<))
16 import Linear hiding (trace, distance)
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.|>)
23 data Amphipod = A | B | C | D deriving (Show, Read, Eq, Ord, Enum)
25 singleStepCost :: Amphipod -> Int
28 singleStepCost C = 100
29 singleStepCost D = 1000
31 type Coord = V2 Int -- r, c
32 _r :: Lens' (V2 Int) Int
34 _c :: Lens' (V2 Int) Int
38 { _destination :: Coord
40 , _transits :: S.Set Coord
41 , _entryRequirement :: Maybe Amphipod
42 } deriving (Show, Eq, Ord)
45 type Steps = M.Map Coord (S.Set Step)
48 { _possibleSteps :: Steps
49 , _roomColumns :: M.Map Amphipod Int
54 type BurrowContext = Reader Burrow
56 type MoveState = M.Map Coord Amphipod
58 data AppliedMove = AppliedMove
59 { _afterMove :: MoveState
60 , _appliedStep :: Step
62 deriving (Show, Eq, Ord)
63 makeLenses ''AppliedMove
66 Agendum { _current :: MoveState
67 , _trail :: Q.Seq MoveState
73 type Agenda = P.MinPQueue Int Agendum
75 type ExploredStates = S.Set MoveState
80 do text <- readFile "data/advent23.txt"
81 -- let (burrow, initState) = mkBurrow text
88 part1 :: String -> Int
89 part1 text = maybe 0 _cost result
91 (burrow, initState) = mkBurrow text
92 result = runReader (searchBurrow initState) burrow
94 part2 :: String -> Int
95 part2 text = maybe 0 _cost result
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
105 searchBurrow :: MoveState -> BurrowContext (Maybe Agendum)
106 searchBurrow initState =
107 do agenda <- initAgenda initState
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}
116 aStar :: Agenda -> ExploredStates -> BurrowContext (Maybe Agendum)
118 | P.null agenda = return Nothing
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
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)
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
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
148 , _trailCost = incurred
149 , _cost = incurred + predicted
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
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)
173 openRoom :: MoveState -> Step -> Bool
176 | otherwise = M.null roomBlockers
177 where e = step ^. entryRequirement
179 tc = step ^. destination . _c
180 roomBlockers = M.filterWithKey (\(V2 _ ac) a -> a /= je && ac == tc) state
182 highestRow :: S.Set Coord -> Step -> Bool
183 highestRow others step = higherRow `S.notMember` others
184 where higherRow = (step ^. destination) & _r +~ 1
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'
191 estimateCost :: MoveState -> BurrowContext Int
193 do rCols <- asks (^. roomColumns)
194 hRow <- asks (^. hallRow)
195 let amphipodCosts = M.mapWithKey (estimateACost rCols hRow) state
196 return $ sum $ M.elems amphipodCosts
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
203 else (r - hRow) + (abs (c - targetCol)) + 1
205 stepCost :: AppliedMove -> Int
206 stepCost aStep = (singleStepCost amphipod) * (aStep ^. appliedStep . distance)
207 where dest = aStep ^. appliedStep . destination
208 amphipod = (aStep ^. afterMove) ! dest
210 isGoal :: MoveState -> BurrowContext Bool
212 do rCols <- asks (^. roomColumns)
213 let misplaced = M.filterWithKey (inWrongRoom rCols) state
214 return $ M.null misplaced
216 inWrongRoom :: M.Map Amphipod Int -> Coord -> Amphipod -> Bool
217 inWrongRoom rCols (V2 _ c) amphipod = c /= rightCol
218 where rightCol = rCols ! amphipod
220 ------------------------------
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
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
239 mkRooms :: [String] -> M.Map Coord Amphipod
240 mkRooms text = M.unions rooms
241 where rooms = map mkRoom $ zip [2..] text
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]
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
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
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
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
272 mkRoomHallRoute :: S.Set Coord -> Coord -> Steps -> Steps
273 mkRoomHallRoute halls here routes = S.foldr' (mkRoomHallRoute1 here) routes halls
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
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
287 mkRoomRoomRoute :: Int -> M.Map Coord Amphipod -> Coord -> Steps -> Steps
288 mkRoomRoomRoute hallRow rooms here routes = M.foldrWithKey' (mkRoomRoomRoute1 hallRow here) routes rooms
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
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
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
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)
315 where maxR = length rows - 1
316 maxC = (length $ head rows) - 1
317 isAmphipod c = c `elem` ("ABCD" :: String)