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