Day 16

Y2022.cabal

@@ -49,15 +49,17 @@ library
         Days.D13
         Days.D14
         Days.D15
+        Days.D16
     build-depends:
         bytestring
       , either
+      , hashable
+      , hashtables              ^>=1.3.1
+      , heap                    ^>=1.0.4
       , lens
       , megaparsec              ^>=9.4.0
       , mtl
       , parser-combinators
-      , hashable
-      , hashtables              ^>=1.3.1
       , text
       , vector                  ^>=0.13.0.0
 

app/Main.hs

@@ -23,6 +23,7 @@ paths =
     , "./data/13.in"
     , "./data/14.in"
     , "./data/15.in"
+    , "./data/16.in"
     ]
 
 solutions :: [(Int, Day, FilePath)]

bench/Bench.hs

@@ -22,6 +22,7 @@ paths =
     , "./data/13.in"
     , "./data/14.in"
     , "./data/15.in"
+    , "./data/16.in"
     ]
 
 solutions :: [(Integer, [Day], FilePath)]

data/16.in

@@ -0,0 +1,51 @@
+Valve NV has flow rate=5; tunnels lead to valves ZV, CG, YB, HX, OY
+Valve NU has flow rate=6; tunnels lead to valves DA, MA, OA, DK
+Valve VU has flow rate=0; tunnels lead to valves PS, FX
+Valve JW has flow rate=0; tunnels lead to valves AA, MD
+Valve RI has flow rate=0; tunnels lead to valves OY, DG
+Valve DG has flow rate=9; tunnels lead to valves TG, RI, DF, EV, KW
+Valve PH has flow rate=7; tunnels lead to valves KW, OW, LT, LZ
+Valve KZ has flow rate=12; tunnels lead to valves ET, QV, CK, MS
+Valve IX has flow rate=0; tunnels lead to valves TS, DO
+Valve MS has flow rate=0; tunnels lead to valves LZ, KZ
+Valve IL has flow rate=0; tunnels lead to valves DO, ET
+Valve EJ has flow rate=20; tunnels lead to valves AV, JY
+Valve DK has flow rate=0; tunnels lead to valves NU, CG
+Valve YB has flow rate=0; tunnels lead to valves NV, PS
+Valve OA has flow rate=0; tunnels lead to valves YA, NU
+Valve DA has flow rate=0; tunnels lead to valves NU, RG
+Valve KO has flow rate=0; tunnels lead to valves AA, TG
+Valve RG has flow rate=4; tunnels lead to valves DF, DA, ZV, MD, LB
+Valve MA has flow rate=0; tunnels lead to valves AA, NU
+Valve OW has flow rate=0; tunnels lead to valves DO, PH
+Valve KW has flow rate=0; tunnels lead to valves DG, PH
+Valve DO has flow rate=14; tunnels lead to valves IX, IL, CZ, OW
+Valve DF has flow rate=0; tunnels lead to valves RG, DG
+Valve TG has flow rate=0; tunnels lead to valves DG, KO
+Valve LB has flow rate=0; tunnels lead to valves RG, FX
+Valve HX has flow rate=0; tunnels lead to valves AA, NV
+Valve GB has flow rate=0; tunnels lead to valves AV, XK
+Valve CG has flow rate=0; tunnels lead to valves DK, NV
+Valve LT has flow rate=0; tunnels lead to valves AO, PH
+Valve FX has flow rate=23; tunnels lead to valves LB, HY, VU
+Valve ET has flow rate=0; tunnels lead to valves IL, KZ
+Valve CK has flow rate=0; tunnels lead to valves UX, KZ
+Valve LZ has flow rate=0; tunnels lead to valves PH, MS
+Valve YA has flow rate=17; tunnels lead to valves JY, OA
+Valve TS has flow rate=0; tunnels lead to valves NO, IX
+Valve NO has flow rate=8; tunnel leads to valve TS
+Valve XK has flow rate=24; tunnel leads to valve GB
+Valve PS has flow rate=18; tunnels lead to valves EV, VU, YB
+Valve AA has flow rate=0; tunnels lead to valves JW, HX, MA, KO
+Valve MD has flow rate=0; tunnels lead to valves JW, RG
+Valve JM has flow rate=19; tunnels lead to valves QV, HY, AO
+Valve AV has flow rate=0; tunnels lead to valves EJ, GB
+Valve AO has flow rate=0; tunnels lead to valves JM, LT
+Valve JY has flow rate=0; tunnels lead to valves YA, EJ
+Valve OY has flow rate=0; tunnels lead to valves NV, RI
+Valve UX has flow rate=13; tunnels lead to valves CZ, CK
+Valve HY has flow rate=0; tunnels lead to valves JM, FX
+Valve EV has flow rate=0; tunnels lead to valves PS, DG
+Valve CZ has flow rate=0; tunnels lead to valves UX, DO
+Valve ZV has flow rate=0; tunnels lead to valves NV, RG
+Valve QV has flow rate=0; tunnels lead to valves JM, KZ

src/Days/D16.hs

@@ -0,0 +1,214 @@
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+
+module Days.D16 where
+
+import Common
+import Data.Char (isAsciiUpper)
+import Data.List (foldl', sortOn)
+import Data.Map (Map)
+import Data.Map qualified as Map
+import Data.Maybe (fromJust, fromMaybe, isJust)
+import Data.Ord (Down (..))
+import Data.Set (Set)
+import Data.Set qualified as Set
+import Data.Text (Text)
+import Data.Tuple (swap)
+import Data.Vector (Vector)
+import Data.Vector qualified as Vector
+import Parse
+import Text.Megaparsec (sepBy, takeWhile1P, (<|>))
+import Text.Megaparsec.Char (space)
+import Text.Megaparsec.Char.Lexer qualified as Lexer
+
+type ValveName = Text
+
+parser :: Parser (Map ValveName Int, Map ValveName [ValveName])
+parser = makeMaps <$> someLines nodeDescription
+  where
+    symbol = Lexer.symbol space
+    lexeme = Lexer.lexeme space
+
+    valve :: Parser Text
+    valve = takeWhile1P Nothing isAsciiUpper
+
+    nodeDescription =
+      (,,)
+        <$ symbol "Valve"
+        <*> lexeme valve
+        <* symbol "has flow rate="
+        <*> number
+        <* symbol ";"
+        <* ( symbol "tunnels lead to valves"
+               <|> symbol "tunnel leads to valve"
+           )
+        <*> sepBy valve (symbol ",")
+
+    makeMaps :: [(Text, Int, [Text])] -> (Map Text Int, Map Text [Text])
+    makeMaps intermediate =
+      let (valve, rate, connections) = unzip3 intermediate
+       in ( Map.fromList (zip valve rate),
+            Map.fromList (zip valve connections)
+          )
+
+floydWarshall :: (Ord a) => Map a [a] -> Map (a, a) Int
+floydWarshall m =
+  let keys = Map.keys m
+   in foldl'
+        ( \acc (k, i, j) ->
+            Map.alter
+              ( let newPath =
+                      do
+                        dist_i_k <- Map.lookup (i, k) acc
+                        dist_k_j <- Map.lookup (k, j) acc
+                        return $ dist_i_k + dist_k_j
+                 in \case
+                      Nothing -> newPath
+                      Just v -> Just $ case newPath of
+                        Nothing -> v
+                        Just newPath_ -> min v newPath_
+              )
+              (i, j)
+              acc
+        )
+        (Map.fromList [((a, b), 1) | a <- keys, b <- fromMaybe [] $ Map.lookup a m])
+        [(k, i, j) | k <- keys, i <- keys, j <- keys]
+
+data Network = Network
+  { rates :: Map ValveName Int,
+    nextBest :: Map ValveName (Vector (ValveName, Int))
+  }
+
+preProcess ::
+  Map ValveName Int ->
+  Map ValveName [ValveName] ->
+  Network
+preProcess rates network =
+  let distances = floydWarshall network
+      possibleStops = Map.filterWithKey (\k v -> v > 0 || k == "AA") rates
+      possibleStopsV = Vector.fromList . Map.toList $ possibleStops
+      nextBest =
+        Map.fromList
+          [ ( from,
+              Vector.fromList
+                . map snd
+                . sortOn (Down . fst)
+                $ [ (rate, (to, dist + 1))
+                    | (to, rate) <- Vector.toList possibleStopsV,
+                      rate > 0,
+                      let maybeDist = Map.lookup (from, to) distances,
+                      isJust maybeDist,
+                      let dist = fromJust maybeDist
+                  ]
+            )
+            | (from, _) <- Vector.toList possibleStopsV
+          ]
+   in Network
+        { rates = possibleStops,
+          nextBest = nextBest
+        }
+
+data SearchState = SearchState
+  { pressure :: Int,
+    opened :: Set ValveName,
+    pos :: (ValveName, ValveName),
+    time :: (Int, Int)
+  }
+
+maximizePressure ::
+  Int ->
+  Bool ->
+  ValveName ->
+  Network ->
+  Int
+maximizePressure deadline withElephant start (Network {..}) =
+  go
+    0
+    SearchState
+      { pressure = 0,
+        opened = Set.empty,
+        pos = (start, start),
+        time = (deadline, if withElephant then deadline else 0)
+      }
+  where
+    go :: Int -> SearchState -> Int
+    go best searcher =
+      let nextStates = branch searcher
+       in if null nextStates
+            then pressure searcher
+            else
+              foldl'
+                ( \best_ nextSearcher ->
+                    if bound nextSearcher > best_
+                      then max (go best_ nextSearcher) best_
+                      else best_
+                )
+                best
+                nextStates
+
+    branch :: SearchState -> [SearchState]
+    branch SearchState {..} =
+      [ let f = if nextTime < snd time then swap else id
+         in SearchState
+              { pressure = pressure + rate * nextTime,
+                opened = Set.insert nextPos opened,
+                pos = f (nextPos, snd pos),
+                time = f (nextTime, snd time)
+              }
+        | (nextPos, dist) <-
+            maybe [] Vector.toList (Map.lookup (fst pos) nextBest),
+          fst time > dist,
+          not $ Set.member nextPos opened,
+          let nextTime = fst time - dist,
+          let rate = fromJust (Map.lookup nextPos rates)
+      ]
+
+    heuristic :: Vector (Vector (ValveName, Int, Int))
+    heuristic =
+      Vector.fromList
+        [ Vector.reverse
+            . Vector.fromList
+            . sortOn (\(_, min_dist, rate) -> (time - min_dist) * rate)
+            $ [ (valve, min_dist, rate)
+                | (valve, rate) <- Map.toList rates,
+                  let min_dist =
+                        minimum
+                          . Vector.map snd
+                          . fromJust
+                          $ Map.lookup valve nextBest,
+                  time > min_dist
+              ]
+          | time <- [0 .. deadline]
+        ]
+
+    bound :: SearchState -> Int
+    bound SearchState {..} = uncurry (go opened) time pressure
+      where
+        go opened max_t min_t bound =
+          case Vector.toList
+            . Vector.filter (\(valve, _, _) -> not $ Set.member valve opened)
+            $ (heuristic Vector.! max_t) of
+            [] -> bound
+            ((valveId, min_dist, flow) : _) ->
+              let max_t' = max_t - min_dist
+                  bound' = bound + flow * max_t'
+                  min_t' = min min_t max_t'
+                  max_t'' = max min_t max_t'
+                  opened' = Set.insert valveId opened
+               in go opened' max_t'' min_t' bound'
+
+part1 :: Network -> Int
+part1 = maximizePressure 30 False "AA"
+
+part2 :: Network -> Int
+part2 = maximizePressure 26 True "AA"
+
+day :: Day
+day =
+  parsecDay
+    (uncurry preProcess <$> parser)
+    ( definitive . part1,
+      definitive . part2
+    )

src/Lib.hs

@@ -27,6 +27,7 @@ import Days.D12 qualified as D12
 import Days.D13 qualified as D13
 import Days.D14 qualified as D14
 import Days.D15 qualified as D15
+import Days.D16 qualified as D16
 
 import Data.ByteString.Char8 qualified as BS
 import Data.Text.IO qualified as T
@@ -54,4 +55,5 @@ days =
     , [D13.day]
     , [D14.day]
     , [D15.day]
+    , [D16.day]
     ]