advent16/MainSPar.hs

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