{-# Language BangPatterns #-}
{-|
Module      : Advent.MinCut
Description : Minimum cut graph algorithm
Copyright   : (c) Eric Mertens, 2023
License     : ISC
Maintainer  : emertens@gmail.com

Karger-Stein approximation of the minimum cut of a graph.

-}
module Advent.MinCut (minCutApprox) where

import Data.Graph.Inductive (match, Gr, insert, edges, order)
import System.Random (StdGen, RandomGen, randomR)

-- | Generate a lazy list of minimun cut approximations.
-- The nodes of the resulting graph will represent the
-- merged components and the remaining edges are the cut.
--
-- The graph is treated as undirected and with uniform edge weight.
minCutApprox ::
  RandomGen gen => Semigroup node =>
  Gr node edge -> gen -> [Gr node edge]
minCutApprox :: forall gen node edge.
(RandomGen gen, Semigroup node) =>
Gr node edge -> gen -> [Gr node edge]
minCutApprox Gr node edge
gr = (gen -> [Gr node edge]) -> Gr node edge -> gen -> [Gr node edge]
forall gen node edge.
(RandomGen gen, Semigroup node) =>
(gen -> [Gr node edge]) -> Gr node edge -> gen -> [Gr node edge]
fastmincut (Gr node edge -> gen -> [Gr node edge]
forall gen node edge.
(RandomGen gen, Semigroup node) =>
Gr node edge -> gen -> [Gr node edge]
minCutApprox Gr node edge
gr) Gr node edge
gr
{-# SPECIALIZE minCutApprox :: Semigroup node => Gr node edge -> StdGen -> [Gr node edge] #-}

-- 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 :: forall gen node edge.
(RandomGen gen, Semigroup node) =>
(gen -> [Gr node edge]) -> Gr node edge -> gen -> [Gr node edge]
fastmincut gen -> [Gr node edge]
k Gr node edge
gr gen
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
6, (Gr node edge
gr', gen
gen') <- Int -> Gr node edge -> gen -> (Gr node edge, gen)
forall gen node edge.
(RandomGen gen, Semigroup node) =>
Int -> Gr node edge -> gen -> (Gr node edge, gen)
contract Int
2 Gr node edge
gr gen
gen = Gr node edge
gr' Gr node edge -> [Gr node edge] -> [Gr node edge]
forall a. a -> [a] -> [a]
: gen -> [Gr node edge]
k gen
gen'
  | Bool
otherwise = (gen -> [Gr node edge]) -> gen -> [Gr node edge]
forall {gen}.
RandomGen gen =>
(gen -> [Gr node edge]) -> gen -> [Gr node edge]
rec ((gen -> [Gr node edge]) -> gen -> [Gr node edge]
forall {gen}.
RandomGen gen =>
(gen -> [Gr node edge]) -> gen -> [Gr node edge]
rec gen -> [Gr node edge]
k) gen
gen -- try twice
  where
    n :: Int
n   = Gr node edge -> Int
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int
order Gr node edge
gr
    t :: Int
t   = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 :: Double)
    rec :: (gen -> [Gr node edge]) -> gen -> [Gr node edge]
rec gen -> [Gr node edge]
k' gen
gen1 | (Gr node edge
gr', gen
gen2) <- Int -> Gr node edge -> gen -> (Gr node edge, gen)
forall gen node edge.
(RandomGen gen, Semigroup node) =>
Int -> Gr node edge -> gen -> (Gr node edge, gen)
contract Int
t Gr node edge
gr gen
gen1 = (gen -> [Gr node edge]) -> Gr node edge -> gen -> [Gr node edge]
forall gen node edge.
(RandomGen gen, Semigroup node) =>
(gen -> [Gr node edge]) -> Gr node edge -> gen -> [Gr node edge]
fastmincut gen -> [Gr node edge]
k' Gr node edge
gr' gen
gen2

-- Karger's algorithm parameterized by vertex stop count
contract ::
  RandomGen gen => Semigroup node =>
  Int -> Gr node edge -> gen -> (Gr node edge, gen)
contract :: forall gen node edge.
(RandomGen gen, Semigroup node) =>
Int -> Gr node edge -> gen -> (Gr node edge, gen)
contract Int
t Gr node edge
gr gen
gen
  | Gr node edge -> Int
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int
order Gr node edge
gr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
t
  , ((Int
l, Int
r), gen
gen1)                <- [Edge] -> gen -> (Edge, gen)
forall gen a. RandomGen gen => [a] -> gen -> (a, gen)
pick (Gr node edge -> [Edge]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [Edge]
edges Gr node edge
gr) gen
gen
  , (Just (Adj edge
li, Int
_, !node
szl, Adj edge
lo), Gr node edge
gr1) <- Int -> Gr node edge -> (MContext node edge, Gr node edge)
forall a b. Int -> Gr a b -> Decomp Gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
Int -> gr a b -> Decomp gr a b
match Int
l Gr node edge
gr
  , (Just (Adj edge
ri, Int
_, !node
szr, Adj edge
ro), Gr node edge
gr2) <- Int -> Gr node edge -> (MContext node edge, Gr node edge)
forall a b. Int -> Gr a b -> Decomp Gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
Int -> gr a b -> Decomp gr a b
match Int
r Gr node edge
gr1
  , let adj :: Adj edge
adj = [(edge, Int)
a | (edge, Int)
a <- Adj edge
li Adj edge -> Adj edge -> Adj edge
forall a. [a] -> [a] -> [a]
++ Adj edge
lo, (edge, Int) -> Int
forall a b. (a, b) -> b
snd (edge, Int)
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r] Adj edge -> Adj edge -> Adj edge
forall a. [a] -> [a] -> [a]
++ Adj edge
ri Adj edge -> Adj edge -> Adj edge
forall a. [a] -> [a] -> [a]
++ Adj edge
ro
  , let gr3 :: Gr node edge
gr3 = (Adj edge, Int, node, Adj edge) -> Gr node edge -> Gr node edge
forall (gr :: * -> * -> *) a b.
DynGraph gr =>
Context a b -> gr a b -> gr a b
insert ([], Int
l, node
szl node -> node -> node
forall a. Semigroup a => a -> a -> a
<> node
szr, Adj edge
adj) Gr node edge
gr2
  = Int -> Gr node edge -> gen -> (Gr node edge, gen)
forall gen node edge.
(RandomGen gen, Semigroup node) =>
Int -> Gr node edge -> gen -> (Gr node edge, gen)
contract Int
t Gr node edge
gr3 gen
gen1

  | Bool
otherwise = (Gr node edge
gr, gen
gen)

-- Selet a random element from a list
pick :: RandomGen gen => [a] -> gen -> (a, gen)
pick :: forall gen a. RandomGen gen => [a] -> gen -> (a, gen)
pick [a]
xs gen
gen =
  case Edge -> gen -> (Int, gen)
forall g. RandomGen g => Edge -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
0, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) gen
gen of
    (Int
i, gen
gen') -> ([a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i, gen
gen')