Updated blog link in comments
[advent-of-code-22.git] / advent16 / MainSubsets.hs
1 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
2 -- Follow up at https://work.njae.me.uk/2023/07/21/optimising-haskell-example-3/
3
4 import Debug.Trace
5
6 import AoC
7 import Data.Text (Text)
8 import qualified Data.Text.IO as TIO
9 import Data.Attoparsec.Text hiding (take, D)
10 import Control.Applicative
11 import qualified Data.PQueue.Prio.Max as P
12 import qualified Data.Set as S
13 import qualified Data.Sequence as Q
14 import qualified Data.Map.Strict as M
15 import Data.Map.Strict ((!))
16 -- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>)))
17 import Data.Sequence ( (|>), Seq((:|>)) )
18 import Data.List
19 -- import Data.List.Split (chunksOf)
20 import Data.Ord
21 import Control.Monad.Reader
22 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
23
24 type RoomID = String
25
26 data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
27 deriving (Eq, Show, Ord)
28 makeLenses ''Tunnel
29
30 data Room = Room
31 { _flowRate :: Int
32 , _tunnels :: S.Set Tunnel
33 } deriving (Eq, Show, Ord)
34 makeLenses ''Room
35
36 type Cave = M.Map RoomID Room
37 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
38
39 type CaveContext = Reader TimedCave
40
41 data SearchState = SearchState
42 { _currentRoom :: RoomID
43 , _currentTime :: Int
44 , _openValves :: S.Set RoomID
45 } deriving (Eq, Show, Ord)
46 makeLenses ''SearchState
47
48 data Agendum =
49 Agendum { _current :: SearchState
50 , _trail :: Q.Seq SearchState
51 , _trailBenefit :: Int
52 , _benefit :: Int
53 } deriving (Show, Eq, Ord)
54 makeLenses ''Agendum
55
56 type Agenda = P.MaxPQueue Int Agendum
57
58 -- state, total flowed so far
59 type ExploredStates = S.Set (SearchState, Int)
60
61 type PartSolutions = M.Map (S.Set RoomID) Int
62
63
64 main :: IO ()
65 main =
66 do dataFileName <- getDataFileName
67 text <- TIO.readFile dataFileName
68 let expandedCave = successfulParse text
69 -- print cave
70 -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
71 -- print $ compress cave
72 let cave = compress expandedCave
73 print $ part1 cave
74 print $ part2 cave
75
76 part1, part2 :: Cave -> Int
77 -- part1 :: Cave -> Int
78 part1 cave = runSearch 30 cave
79 part2 cave = maximum combinations
80 where rawSolutions = runSearchAll 26 cave
81 solutionList = M.toList rawSolutions
82 combinations = [ (f1 + f2)
83 | (p, f1) <- solutionList
84 , (e, f2) <- solutionList
85 , p < e
86 , S.disjoint p e
87 ]
88
89 includeAgendum :: PartSolutions -> Agendum -> CaveContext PartSolutions
90 includeAgendum results agendum =
91 do cf <- currentFlow (agendum ^. current)
92 timeLimit <- asks getTimeLimit
93 let timeLeft = timeLimit - timeOf (agendum ^. current)
94 let remainingFlow = cf * timeLeft
95 let totalFlow = remainingFlow + agendum ^. trailBenefit
96 let visitedSet = agendum ^. current . openValves
97 let currentBest = M.findWithDefault 0 visitedSet results
98 if totalFlow > currentBest
99 then return (M.insert visitedSet totalFlow results)
100 else return results
101
102 runSearch :: Int -> Cave -> Int
103 runSearch timeLimit cave = maybe 0 _benefit result
104 where result = runReader (searchCave "AA") (TimedCave cave timeLimit sortedRooms)
105 sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
106
107 runSearchAll :: Int -> Cave -> PartSolutions
108 runSearchAll timeLimit cave = result
109 where result = runReader (searchCaveAll "AA") (TimedCave cave timeLimit sortedRooms)
110 sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
111
112
113 searchCave :: String -> CaveContext (Maybe Agendum)
114 searchCave startRoom =
115 do agenda <- initAgenda startRoom
116 aStar agenda S.empty
117
118 searchCaveAll :: String -> CaveContext PartSolutions
119 searchCaveAll startRoom =
120 do agenda <- initAgenda startRoom
121 allSolutions agenda S.empty M.empty
122
123 initAgenda :: String -> CaveContext Agenda
124 initAgenda startID =
125 do let startState = emptySearchState startID
126 b <- estimateBenefit startState
127 return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
128
129 aStar :: Agenda -> ExploredStates -> CaveContext (Maybe Agendum)
130 aStar agenda closed
131 -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
132 -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
133 -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
134 -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
135 | P.null agenda = return Nothing
136 | otherwise =
137 do let (_, currentAgendum) = P.findMax agenda
138 let reached = currentAgendum ^. current
139 nexts <- candidates currentAgendum closed
140 let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
141 reachedGoal <- isGoal currentAgendum
142 let cl = (reached, currentAgendum ^. trailBenefit)
143 if reachedGoal
144 then return (Just currentAgendum)
145 else if (cl `S.member` closed)
146 then aStar (P.deleteMax agenda) closed
147 else aStar newAgenda (S.insert cl closed)
148
149 allSolutions :: Agenda -> ExploredStates -> PartSolutions -> CaveContext PartSolutions
150 allSolutions agenda closed foundSolutions
151 | P.null agenda = return foundSolutions
152 | otherwise =
153 do let (_, currentAgendum) = P.findMax agenda
154 let reached = currentAgendum ^. current
155 nexts <- candidates currentAgendum closed
156 let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
157 reachedGoal <- isGoal currentAgendum
158 let cl = (reached, currentAgendum ^. trailBenefit)
159 newFoundSolutions <- includeAgendum foundSolutions currentAgendum
160 if reachedGoal
161 then allSolutions (P.deleteMax agenda) closed newFoundSolutions
162 else if (cl `S.member` closed)
163 then allSolutions (P.deleteMax agenda) closed foundSolutions
164 else allSolutions newAgenda (S.insert cl closed) newFoundSolutions
165
166
167 candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
168 candidates agendum closed =
169 do let candidate = agendum ^. current
170 let previous = agendum ^. trail
171 let prevBenefit = agendum ^. trailBenefit
172 succs <- successors candidate
173 succAgs <- mapM (makeAgendum previous prevBenefit) succs
174 let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
175 return nonloops
176
177 emptySearchState :: RoomID -> SearchState
178 emptySearchState startID = SearchState
179 { _currentRoom = startID
180 , _currentTime = 0
181 , _openValves = S.empty
182 }
183
184 currentFlow :: SearchState -> CaveContext Int
185 currentFlow state =
186 do cave <- asks getCave
187 let valves = state ^. openValves
188 let presentRooms = cave `M.restrictKeys` valves
189 return $ sumOf (folded . flowRate) presentRooms
190
191 timeOf :: SearchState -> Int
192 timeOf state = state ^. currentTime
193
194 successors :: SearchState -> CaveContext (Q.Seq SearchState)
195 successors state =
196 do isFF <- isFullFlow state
197 cave <- asks getCave
198 timeLimit <- asks getTimeLimit
199 let here = state ^. currentRoom
200 let opened = state ^. openValves
201 let now = state ^. currentTime
202 let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
203 let moves = [ SearchState
204 { _currentRoom = (t ^. tunnelTo)
205 , _currentTime = now + (t ^. tunnelLength)
206 , _openValves = opened
207 }
208 | t <- remaining
209 , now + (t ^. tunnelLength) <= timeLimit
210 ]
211 let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 && now < timeLimit
212 then [SearchState { _currentRoom = here, _currentTime = now + 1, _openValves = S.insert here opened }]
213 else []
214 let nexts = if null opens then moves else opens
215 let nexts' = if null nexts
216 then [ SearchState
217 { _currentRoom = here
218 , _currentTime = timeLimit
219 , _openValves = opened
220 } ]
221 else nexts
222 let succs = Q.fromList nexts'
223 if isFF || (Q.null succs)
224 then return $ Q.singleton (state & currentTime .~ timeLimit)
225 else return succs
226
227 estimateBenefit :: SearchState -> CaveContext Int
228 estimateBenefit here =
229 do cave <- asks getCave
230 timeLimit <- asks getTimeLimit
231 let timeRemaining = timeLimit - (timeOf here)
232 cf <- currentFlow here
233 sortedValves <- asks getSortedRooms
234 let opened = here ^. openValves
235 let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
236 let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
237 return $ (cf * timeRemaining) + otherValveFlows
238
239 makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> CaveContext Agendum
240 makeAgendum previous prevBenefit newState =
241 do predicted <- estimateBenefit newState -- (Q.length previous)
242 -- cf <- currentFlow newState
243 oldFlow <- lastFlow previous (timeOf newState)
244 let newTrail = previous |> newState
245 let incurred = prevBenefit + oldFlow
246 return Agendum { _current = newState
247 , _trail = newTrail
248 , _trailBenefit = incurred
249 , _benefit = incurred + predicted
250 }
251
252 lastFlow :: Q.Seq SearchState -> Int -> CaveContext Int
253 lastFlow Q.Empty _ = return 0
254 lastFlow (_ :|> previous) newTime =
255 do cf <- currentFlow previous
256 let dt = newTime - (timeOf previous)
257 return (cf * dt)
258
259 isGoal :: Agendum -> CaveContext Bool
260 isGoal agendum =
261 do timeLimit <- asks getTimeLimit
262 let s = agendum ^. current
263 return $ (timeOf s) == timeLimit
264
265 isFullFlow :: SearchState -> CaveContext Bool
266 isFullFlow state =
267 do cave <- asks getCave
268 cf <- currentFlow state
269 let ff = sumOf (folded . flowRate) cave
270 return (cf == ff)
271
272 compress :: Cave -> Cave
273 compress cave = M.mapWithKey (compressRoom cave) cave
274
275 compressRoom :: Cave -> RoomID -> Room -> Room
276 compressRoom cave here room = room & tunnels .~ t'
277 where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
278
279 reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
280 reachableFrom _ [] _ routes = routes
281 reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
282 | here `S.member` found = reachableFrom cave boundary found routes
283 | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
284 where exits = (cave ! here) ^. tunnels
285 exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
286 legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
287 routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
288 then routes
289 else S.insert tunnel routes
290
291 -- Parse the input file
292
293 caveP :: Parser Cave
294 valveP :: Parser (RoomID, Room)
295 roomP :: Parser Room
296 tunnelsP :: Parser (S.Set Tunnel)
297 tunnelTextP :: Parser Text
298
299 caveP = M.fromList <$> valveP `sepBy` endOfLine
300 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
301 roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
302 -- where roomify v ts = Room {flowRate = v, tunnels = ts }
303 tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
304 tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
305
306 successfulParse :: Text -> Cave
307 successfulParse input =
308 case parseOnly caveP input of
309 Left _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
310 Right cave -> cave