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