2 files changed, 39 insertions(+), 46 deletions(-)

M 2022/16.hs
M lib/Qulude.hs

@@ 25,9 25,9 @@ indexed = go 0
   where go _ [] = []
         go i (x:xs) = (i,x):go (succ i) xs
 
-data Survey = Survey { names :: Array Int Name
-                     , rates :: Array Int Int
-                     , tunnels :: Array Int [Int]
+data Survey = Survey { rates :: Array Int Int
+                     , distances :: Map (Int,Int) Int
+                     , initiallyClosedValves :: [ Int ]
                      }
   deriving (Show, Eq)
 


@@ 44,71 44,59 @@ parser = mk <$> many valve
           newline
           return $ (v, r, ds)
 
-mk vs = Survey ns rs ts
+mk vs = Survey { rates, distances, initiallyClosedValves }
   where is = indexed vs
         n = length is
-        ns = Array.array (0,pred n) $ flip fmap is $ \(i, (n, _, _)) -> (i,n)
-        rs = Array.array (0,pred n) $ flip fmap is $ \(i, (_, r, ds)) -> (i,r)
-        ts = Array.array (0,pred n) $ flip fmap is $ \(i, (_, _, ds)) -> (i, g 0 <$> ds)
-        g i d | d == ns ! i = i
+        r = (0, pred $ length is)
+        names = Array.array r $ flip fmap is $ \(i, (n, _, _)) -> (i,n)
+        tunnels = Array.array r $ flip fmap is $ \(i, (_, _, ds)) -> (i, g 0 <$> ds)
+        rates = Array.array r $ flip fmap is $ \(i, (_, r, ds)) -> (i,r)
+        initiallyClosedValves = Maybe.mapMaybe (\(i, (_, r, _)) -> if r > 0 then Just i else Nothing) is
+        g i d | d == names ! i = i
         g i d | otherwise = g (succ i) d
+        distances = Map.fromList [ ((i, j), d i j) | i <- 0:initiallyClosedValves, j <- initiallyClosedValves, i /= j ]
+        ns = Set.fromList $ Array.range r
+        d i j = dijkstra i j ns (\k -> fmap (flip (,) 1) (tunnels ! k))
 
 data Action = Noop | Move Int | Open
   deriving ( Show, Eq )
 
 data St = St { loc :: Int
              , pressure :: Int
-             , valves :: Array Int Bool
+             , openValves :: [ Int ]
+             , closedValves :: [ Int ]
              , time :: Int
              }
   deriving ( Show, Eq )
 
 initSt s = St { loc = 0
               , pressure = 0
-              , valves = Array.listArray (Array.bounds $ names s) $ repeat False
+              , openValves = []
+              , closedValves = initiallyClosedValves s
               , time = 0
               }
 
-ix (Survey { names }) = Array.indices names
+releasePressure s@(Survey { rates }) (St { openValves }) t = sum $ fmap (\i -> t * rates ! i) openValves
 
-releasePressure s@(Survey { rates }) (St { valves }) t = sum $ fmap f (ix s)
-  where f i | valves ! i = t * (rates ! i)
-        f i | otherwise = 0
+move s i st = st { time = time st + d
+                 , loc = i
+                 , pressure = pressure st + releasePressure s st d
+                 }
+  where d = (Map.!) (distances s) (loc st, i)
 
-nextValves s@(Survey { rates }) (St { loc, valves }) = catMaybes $ fmap f (ix s)
-  where f i | i == loc = Nothing
-        f i | otherwise  = if (valves ! i) || (rates ! i /= 0) then Just i else Nothing
+open s i st = st { time = time st + 1
+                 , pressure = pressure st + releasePressure s st 1
+                 , closedValves = List.delete i $ closedValves st
+                 , openValves = i:openValves st
+                 }
 
-move s ds i st = st { time = time st + d
-                    , loc = i
-                    , pressure = pressure st + releasePressure s st d
-                    }
-  where d = (Map.!) ds (loc st, i)
-
-open s ds i st = st { time = time st + 1
-                    , pressure = pressure st + releasePressure s st 1
-                    , valves = valves st // [ (i, True) ]
-                    }
-
-mapMaybe :: (a -> Maybe b) -> [a] -> [b]
-mapMaybe f [] = []
-mapMaybe f (a:as) = case f a of
-                      Just b -> b:(mapMaybe f as)
-                      Nothing -> mapMaybe f as
-
-next s ds st@(St { loc, pressure, time, valves }) = mapMaybe f $ nextValves s st
-  where f i = let d = (Map.!) ds (loc, i) in
-              if 30 - time > d then Just (open s ds i $ move s ds i st) else Nothing
-
-distances s = Map.fromList [ ((i, j), d i j) | i <- ix s, j <- ix s, i /= j ]
-  where vs = length $ nextValves s $ initSt s
-        ns = Set.fromList (ix s)
-        d i j = dijkstra i j ns (\k -> fmap (flip (,) 1) (tunnels s ! k))
+next s st@(St { loc, pressure, time, closedValves }) = Maybe.mapMaybe f closedValves
+  where f i = let d = (Map.!) (distances s) (loc, i) in
+              if 30 - time > d + 1 then Just (open s i $ move s i st) else Nothing
 
 partA s = go 0 [ initSt s ]
-  where ds = distances s
-        go !m [] = m
-        go !m (st:acc) = case next s ds st of
+  where go !m [] = m
+        go !m (st:acc) = case next s st of
                            [] -> let m' = runOutTheClock st in
                                  go (max m m') acc
                            sts -> go m (sts `combine` acc)

@@ 41,11 41,16 @@ module Qulude ( module Prelude
 
               , (\\), (&)
               , (><), (<|), (|>), Seq ( (:<|), (:|>) )
+
               , fromJust
               , fromMaybe
               , catMaybes
+              , mapMaybe
+
               , swap
+
               , fromRight', fromLeft', fromEither
+
               , foldl'
 
               , unsafePerformIO


@@ 61,7 66,7 @@ import Prelude hiding ( lines, cycle )
 
 import Data.List ( (\\), foldl' )
 import Data.Sequence ( (><), (<|), (|>), Seq ( (:<|), (:|>) ) )
-import Data.Maybe ( fromJust, fromMaybe, catMaybes )
+import Data.Maybe ( fromJust, fromMaybe, catMaybes, mapMaybe )
 import Data.Tuple ( swap )
 import Data.Function ( (&) )
 import Data.Either