cleanup.

mostly merge function, I think it flows better using State.

happy with d25 now.

src/Javran/AdventOfCode/Y2023/Day25.hs

@@ -1,6 +1,6 @@
 module Javran.AdventOfCode.Y2023.Day25 () where
 
-import Control.Monad
+import Control.Monad.State.Strict
 import Control.Monad.Writer.CPS
 import Data.Char
 import qualified Data.IntSet as IS
@@ -62,49 +62,79 @@ type G = M.Map V (S.Set V)
 type W = M.Map (MinMax V) Int
 type GW = (G, W)
 
+{-
+
+  `merge s t gw` merges vertices `s` and `t` together in graph `gw`.
+
+  - Value of this new vertex is simply set union of `s` and `t`.
+
+  - Edge weight between `s` and `t` is ignored (if any).
+
+  - Connectivity outside of `s` and `t` are preserved. If they both connect to the same vertex,
+    new edge weight is the sum of old weights.
+
+ -}
 merge :: V -> V -> GW -> GW
-merge s t (g0, w0) = do
+merge s t = execState do
+  (g0, w0) <- get
   let
     st = s <> t
-    sConns = maybe [] S.toList (g0 M.!? s)
-    tConns = maybe [] S.toList (g0 M.!? t)
-    gFin :: M.Map V (S.Set V)
-    gFin = M.insert st (S.delete t $ S.delete s $ S.fromList $ sConns <> tConns) g3
-      where
-        g1 = M.delete t . M.delete s $ g0
-        g2 = foldr (\i -> M.adjust (S.insert st . S.delete s) i) g1 sConns
-        g3 = foldr (\i -> M.adjust (S.insert st . S.delete t) i) g2 tConns
-    wFin :: M.Map (MinMax V) Int
-    wFin = M.unionWith (+) w1 stWeights
-      where
-        w1 = M.filterWithKey (\(MinMax (a, b)) _v -> not (elem a [s, t] || elem b [s, t])) w0
-        stWeights :: M.Map (MinMax V) Int
-        stWeights = M.fromListWith (+) do
-          (x, conns) <- [(s, filter (/= t) sConns), (t, filter (/= s) tConns)]
-          y <- conns
-          Just v <- [w0 M.!? minMaxFromPair (x, y)]
-          pure (minMaxFromPair (st, y), v)
-
-  (gFin, wFin)
-
-minimumCutPhase :: G -> W -> V -> Maybe ((V, V), Int)
-minimumCutPhase g w a = do
-  -- initialize
+    sConns = fromMaybe S.empty (g0 M.!? s)
+    tConns = fromMaybe S.empty (g0 M.!? t)
+    stConns = sConns <> tConns
+
+  -- remove s and t
+  modify $ first $ M.delete t . M.delete s
+  -- replace s or t with st
+  modify $ first $ \g -> foldr (M.adjust (S.insert st . S.delete t . S.delete s)) g stConns
+  -- move connections over to st
+  modify $ first $ M.insert st (S.delete t . S.delete s $ stConns)
+
+  -- remove edges whose one end is s or t
+  modify $ second $ M.filterWithKey (\(MinMax (a, b)) _v -> notElem a [s, t] && notElem b [s, t])
   let
-    q0 :: PQ.PSQ V (Down Int)
-    q0 = PQ.singleton a (Down 0)
+    stWeights :: M.Map (MinMax V) Int
+    stWeights = M.fromListWith (+) do
+      -- all connections except that between s and t.
+      (x, conns) <- [(s, S.delete t sConns), (t, S.delete s tConns)]
+      y <- S.toList conns
+      Just v <- [w0 M.!? minMaxFromPair (x, y)]
+      pure (minMaxFromPair (st, y), v)
+  -- add new weights related to st.
+  modify $ second $ M.unionWith (+) stWeights
 
+{-
+  The "MinimumCutPhase" algorithm as described in the paper.
+
+  This feels familiar to Dijkstra's algorithm: we begin with a singleton
+  set of vertices and expand it to the entire graph.
+
+  Meanwhile a `PSQ` is used as heap to extract maximum value (with `Down`) efficiently
+  to determine which vertex should be added next.
+
+ -}
+minimumCutPhase :: G -> W -> V -> Maybe ((V, V), Int)
+minimumCutPhase g w a =
   fix
     ( \go aSet q acc -> case PQ.minView q of
         Nothing -> case acc of
-          (s, v) : (t, _) : _ -> Just ((s, t), v)
+          (s, v) : (t, _) : _ ->
+            {-
+              As this list is accumulated in reverse order,
+              we extract first 2 elements which should be vertices we need to merge.
+             -}
+            Just ((s, t), v)
           _ -> Nothing
-        Just (z PQ.:-> (Down zW), q1) -> do
-          -- insert z into the set A
+        Just (z PQ.:-> (Down zW), q1) ->
+          -- insert z into `aSet`
           let
             vExtras = do
-              v <- maybe [] S.toList (g M.!? z)
-              guard $ S.notMember v aSet
+              -- collect weights related to z to update the queue.
+              v <-
+                maybe
+                  []
+                  (S.toList . (\s -> S.difference s aSet))
+                  (g M.!? z)
               pure (v, w M.! minMaxFromPair (z, v))
             q2 = foldr upd q1 vExtras
               where
@@ -115,19 +145,24 @@ minimumCutPhase g w a = do
                         Just (Down wOld) -> Just (Down (wOld + wv))
                     )
                     v
-          go (S.insert z aSet) q2 ((z, zW) : acc)
+           in
+            go (S.insert z aSet) q2 ((z, zW) : acc)
     )
     (S.singleton a)
-    q0
+    (PQ.singleton a (Down 0))
     []
 
+{-
+  "MinimumCut" algorithm, see paper.
+ -}
 stoerWagner :: GW -> Maybe (ArgMin Int V)
 stoerWagner p0@(g0, _) = execWriter $ runW p0
   where
     a : _ = M.keys g0
     runW = fix \go cur@(curG, curW) ->
       case minimumCutPhase curG curW a of
         Just ((s, t), w) -> do
+          -- announce this cut to update current best.
           tell (Just (Min (Arg w s)))
           go (merge s t cur)
         Nothing -> pure ()