1 files changed, 46 insertions, 0 deletions
diff --git a/2022/16/solve.hs b/2022/16/solve.hs
new file mode 100644
index 0000000..b7744c2
--- /dev/null
+++ b/2022/16/solve.hs
@@ -0,0 +1,46 @@
+{-# LANGUAGE TupleSections #-}
+
+import           Control.Monad
+import qualified Data.HashMap                  as Map
+import           Data.HashMap                   ( Map )
+import           Data.List
+import           Data.MemoTrie
+import           Text.Regex.TDFA
+
+m !!! k = Map.findWithDefault 9999999 k m
+
+parse input = (valves, Map.fromList rates, Map.fromList dists)
+ where
+  match :: String -> [String]
+  match s = getAllTextMatches $ s =~ "([A-Z]{2}|[0-9]+)"
+  matches = map match input
+  valves  = map head matches
+  rates = [ (valve, read rate :: Int) | (valve : rate : _) <- matches, rate /= "0" ]
+  dists = [ ((valve, pipe), 1) | (valve : _ : pipes) <- matches, pipe <- pipes ]
+
+precompute valves distances = updated distances prod
+ where
+  prod = sequence $ replicate 3 valves
+  dist (x : y : z : _) m = min (m !!! (y, z)) ((m !!! (y, x)) + (m !!! (x, z)))
+  updated = foldr (\c@(_ : y : z : _) m -> Map.insert (y, z) (dist c m) m)
+
+solve' :: Int -> Int -> String -> [String] -> Map String Int -> Map (String, String) Int -> Int
+solve' p time current valves rates dists = maximum $ single ++ case p of
+  1 -> [0]
+  2 -> [solve 1 26 "AA" valves rates dists]
+ where
+  single = [ each valve | valve <- valves, (dists !!! (current, valve)) < time ]
+  new_rate = (rates Map.!)
+  new_time = (+ (-1)) . (-) time . (dists !!!) . (current, ) -- troll
+  next valve = solve p (new_time valve) valve (delete valve valves) rates dists
+  each = ap ((+) . liftM2 (*) new_rate new_time) next -- lol
+
+solve = memo solve'
+
+main :: IO ()
+main = do
+  input <- lines <$> readFile "input"
+  let (valves, rates, dists) = parse input
+      dists'                 = precompute valves dists
+  putStrLn $ show $ solve 1 30 "AA" (Map.keys rates) rates dists'
+  putStrLn $ show $ solve 2 26 "AA" (Map.keys rates) rates dists'