From 22e44a6fa9fae35750b2f21f3112344db920af16 Mon Sep 17 00:00:00 2001
From: Eric Mertens <emertens@galois.com>
Date: Tue, 26 Dec 2023 21:29:42 -0600
Subject: [PATCH] lazy list to externalize stop condition

---
 solutions/src/2023/25.hs | 48 ++++++++++++++++++++--------------------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/solutions/src/2023/25.hs b/solutions/src/2023/25.hs
index d9d7532..ea92352 100644
--- a/solutions/src/2023/25.hs
+++ b/solutions/src/2023/25.hs
@@ -1,4 +1,4 @@
-{-# Language QuasiQuotes, BangPatterns #-}
+{-# Language QuasiQuotes, BangPatterns, TransformListComp, BlockArguments #-}
 {-|
 Module      : Main
 Description : Day 25 solution
@@ -45,49 +45,49 @@ module Main (main) where
 import Advent (format, ordNub)
 import Advent.Tokenize (autoTokenize)
 import Data.Graph.Inductive (Gr, UGr, insert, match, size, labNodes, edges, nmap, mkUGraph, noNodes)
-import System.Random (getStdRandom, randomR, RandomGen, StdGen)
+import System.Random (newStdGen, randomR, RandomGen)
+import Data.Semigroup (Sum(Sum))
 
 main :: IO ()
 main =
  do input <- [format|2023 25 (%s:( %s)*%n)*|]
-    let g = nmap (const 1) (simpleGraph (autoTokenize input))
-        loop gen =
-          case fastmincut g gen of
-            (g', gen')
-              | size g' == 3 -> (g', gen')
-              | otherwise    -> loop gen'
-    g' <- getStdRandom loop
-    print (product [sz | (_, sz) <- labNodes g'])
+    let g = nmap (const (Sum (1::Int))) (simpleGraph (autoTokenize input))
+        makeCandidates = fastmincut makeCandidates g
+    gs <- makeCandidates <$> newStdGen
+    print (product [sz | g' <- gs, 3 == size g', then take 1, (_, Sum sz) <- labNodes g'])
 
 -- Transform the input format into an fgl unlabeled graph
 simpleGraph :: [(Int, [Int])] -> UGr
 simpleGraph input =
   mkUGraph
-    (ordNub [n | (k,vs) <- input, n <- k:vs])
-    [(k,v) | (k,vs) <- input, v <- vs]
+    (ordNub [n | (k, vs) <- input, n <- k : vs])
+    [(k, v) | (k, vs) <- input, v <- vs]
 
--- Karger–Stein algorithm (specialized to find mincuts of size 3)
-fastmincut :: Gr Int e -> StdGen -> (Gr Int e, StdGen)
-fastmincut g gen
-  | n <= 6    = contract 2 g gen
-  | otherwise =
-      case rec gen of
-        (g', gen') | size g' == 3 -> (g', gen')
-                   | otherwise    -> rec gen'
+-- Karger–Stein algorithm parameterized over the continuation
+-- that consumes the random generator. This allows the implementation
+-- to generate an infinite list of candidates to be selected from.
+-- The 'Semigroup' instance is used to combine merged nodes.
+fastmincut ::
+  RandomGen gen => Semigroup node =>
+  (gen -> [Gr node edge]) {- ^ continuation -} ->
+  Gr node edge -> gen -> [Gr node edge]
+fastmincut k g gen
+  | n <= 6, (g', gen') <- contract 2 g gen = g' : k gen'
+  | otherwise = rec (rec k) gen -- try twice
   where
     n   = noNodes g
     t   = ceiling (1 + fromIntegral n / sqrt 2 :: Double)
-    rec = uncurry fastmincut . contract t g
+    rec kr = uncurry (fastmincut kr) . contract t g
 
--- Karger algorithm parameterized by stop condition
-contract :: RandomGen g => Int -> Gr Int e -> g -> (Gr Int e, g)
+-- Karger's algorithm parameterized by vertex stop count
+contract :: (RandomGen g, Semigroup n) => Int -> Gr n e -> g -> (Gr n e, g)
 contract t g gen
   | noNodes g > t
   , ((l, r), gen')               <- pick (edges g) gen
   , (Just (li, _, !szl, lo), g1) <- match l g
   , (Just (ri, _, !szr, ro), g2) <- match r g1
   , let adj = [a | a <- li ++ lo, snd a /= r] ++ ri ++ ro
-  , let g3 = insert ([], l, szl + szr, adj) g2
+  , let g3 = insert ([], l, szl <> szr, adj) g2
   = contract t g3 gen'
 
   | otherwise = (g, gen)