day 16 part2

2022/day16.hs

@@ -5,8 +5,9 @@ import Debug.Trace (traceShowId, traceShow, trace)
 import Data.Maybe (fromMaybe)
 
 import Bfs
-import Data.List (maximumBy, permutations, delete)
+import Data.List (maximumBy, permutations, delete, singleton)
 import Data.Function (on)
+import Control.Monad.Trans.State (State, get, put, evalState)
 
 data Valve = V { flow' :: Int, neighs' :: [String] } deriving Show
 type ZMap = M.Map String Valve
@@ -63,7 +64,7 @@ perms limit m = perms' "AA" limit elems where
     perms' current remaining as = do
         a <- as
         let distance = dists (m M.! current) M.! a
-        if distance < remaining then do
+        [a] : if distance < remaining then do
             let l = delete a as
             ls <- perms' a (remaining-distance) l
             pure $ a:ls
@@ -88,14 +89,54 @@ traceMap cond f (x:y:xs)
     | cond x y = let result = f x in traceShow (result, x) (result: traceMap cond f (y:xs))
     | otherwise = f x : traceMap cond f (y:xs)
 
+type EleMemo = M.Map (S.Set String) Int
+eleLookup :: NzMap -> S.Set String -> State EleMemo Int
+eleLookup nzmap permutation = do
+    mmap <- get
+    case M.lookup permutation mmap of
+        Just x -> pure x
+        Nothing -> do
+            let valves = M.fromList $ filter (\(name, valve) -> S.notMember name permutation) (M.toList nzmap)
+            let elePerms = traceShow ("elephant", permutation) $ perms 26 valves
+            let exec p = execPerm nzmap p 0 "AA" 26
+            let flow = maximum $ map exec elePerms
+            put $ M.insert permutation flow mmap
+            pure flow
+
+part2 :: NzMap -> [String] -> State EleMemo Int
+part2 nzmap permutation = do
+    let me = execPerm nzmap permutation 0 "AA" 26
+    elephant <- eleLookup nzmap (S.fromList permutation)
+    pure $ me + elephant
+    -- case elephant of
+    --     Just ele -> pure $ me + ele
+    --     Nothing -> do
+    --         let valves = M.fromList $ filter (\(name, valve) -> S.notMember name permset) (M.toList nzmap)
+    --         let elePerms = traceShow ("elephant", permutation) $ perms 26 valves
+    --         let exec p = execPerm nzmap p 0 "AA" 26
+    --         let ele = maximum $ map exec elePerms
+    --         put $ M.insert permset ele state
+    --         pure $ me + ele
+
 main :: IO ()
 main = do
     x <- prepass . M.fromList . map parseValve . lines <$> getContents
     mapM_ print x
 
-    let all = perms 30 x
+    -- let all = perms 30 x
-    let exec p = execPerm x p 0 "AA" 30
+    -- let exec p = execPerm x p 0 "AA" 30
-    let result = traceMap (\a b -> take 2 a /= take 2 b) exec all
+    -- let result = traceMap (\a b -> take 2 a /= take 2 b) exec all
-    putStr "part 1: "
+    -- putStr "part 1: "
-    print $ maximum result
+    -- print $ maximum result
     -- print $ fst $ search x S.empty "AA" 30 M.empty
+
+    putStr "part 2: "
+    -- let initialEle = maximumBy (compare `on` fst) $ map (\p -> (execPerm x p 0 "AA" 26, [S.fromList p])) (perms 26 x)
+    -- print initialEle
+
+
+    let all = perms 26 x
+    let action = mapM (part2 x) all
+    let result = evalState action M.empty
+    print $ maximum result
+    pure ()