module HGraph.Directed.Connectivity.Flow
       ( maxFlow
       , maxDisjointPaths
       , minCut
       , minCutI
       )
where

import Data.List
import HGraph.Directed
import qualified Data.Map as M
import qualified Data.Set as S
import Control.Monad

maxFlow :: (Ord a, Adjacency t, DirectedGraph t) => t a -> a -> a -> M.Map (a, a) Bool
maxFlow :: t a -> a -> a -> Map (a, a) Bool
maxFlow t a
d a
s a
t = Map (a, a) Bool -> Map (a, a) Bool
maxFlow' (Map (a, a) Bool -> Map (a, a) Bool)
-> Map (a, a) Bool -> Map (a, a) Bool
forall a b. (a -> b) -> a -> b
$ ((a, a) -> Map (a, a) Bool -> Map (a, a) Bool)
-> Map (a, a) Bool -> [(a, a)] -> Map (a, a) Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(a, a)
a -> (a, a) -> Bool -> Map (a, a) Bool -> Map (a, a) Bool
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (a, a)
a Bool
False) Map (a, a) Bool
forall k a. Map k a
M.empty (t a -> [(a, a)]
forall (t :: * -> *) a. DirectedGraph t => t a -> [(a, a)]
arcs t a
d)
  where
    maxFlow' :: Map (a, a) Bool -> Map (a, a) Bool
maxFlow' Map (a, a) Bool
flow 
      | [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
p = Map (a, a) Bool
flow
      | Bool
otherwise = Map (a, a) Bool -> Map (a, a) Bool
maxFlow' Map (a, a) Bool
flow'
      where
        p :: [a]
p = t a -> a -> a -> Map (a, a) Bool -> [a]
forall a (t :: * -> *).
(Ord a, Adjacency t) =>
t a -> a -> a -> Map (a, a) Bool -> [a]
shortestPathResidual t a
d a
s a
t Map (a, a) Bool
flow
        flow' :: Map (a, a) Bool
flow' = ((a, a) -> Map (a, a) Bool -> Map (a, a) Bool)
-> Map (a, a) Bool -> [(a, a)] -> Map (a, a) Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Bool -> Bool) -> (a, a) -> Map (a, a) Bool -> Map (a, a) Bool
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust Bool -> Bool
not) Map (a, a) Bool
flow ([(a, a)] -> Map (a, a) Bool) -> [(a, a)] -> Map (a, a) Bool
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
p ([a] -> [a]
forall a. [a] -> [a]
tail [a]
p)

shortestPathResidual :: t a -> a -> a -> Map (a, a) Bool -> [a]
shortestPathResidual t a
d a
s a
t Map (a, a) Bool
flow = Set a -> Map a a -> [a]
path (a -> Set a
forall a. a -> Set a
S.singleton a
s) Map a a
forall k a. Map k a
M.empty
  where
    path :: Set a -> Map a a -> [a]
path Set a
active Map a a
preds
      | a
t a -> Map a a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a a
preds = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ Map a a -> a -> [a]
makePath Map a a
preds a
t
      | Set a -> Bool
forall a. Set a -> Bool
S.null Set a
active = []
      | Bool
otherwise = Set a -> Map a a -> [a]
path ([a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ Map a a -> [a]
forall k a. Map k a -> [k]
M.keys Map a a
newPred) (Map a a
preds Map a a -> Map a a -> Map a a
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map a a
newPred)
        where
          newPred :: Map a a
newPred = [(a, a)] -> Map a a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, a)] -> Map a a) -> [(a, a)] -> Map a a
forall a b. (a -> b) -> a -> b
$ [ (a
u,a
v)
                             | a
v <- Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
active
                             , a
u <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
outneighbors t a
d a
v
                             , (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Map (a, a) Bool
flow Map (a, a) Bool -> (a, a) -> Bool
forall k a. Ord k => Map k a -> k -> a
M.! (a
v,a
u)) Bool -> Bool -> Bool
&& (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a
u a -> Map a a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a a
preds)
                             ]
                             [(a, a)] -> [(a, a)] -> [(a, a)]
forall a. [a] -> [a] -> [a]
++
                             [ (a
u,a
v)
                             | a
v <- Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
active
                             , a
u <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
inneighbors t a
d a
v
                             , Map (a, a) Bool
flow Map (a, a) Bool -> (a, a) -> Bool
forall k a. Ord k => Map k a -> k -> a
M.! (a
u, a
v) Bool -> Bool -> Bool
&& (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a
u a -> Map a a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a a
preds)
                             ]
    makePath :: Map a a -> a -> [a]
makePath Map a a
preds a
v
      | a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
s = [a
v]
      | Bool
otherwise = a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Map a a -> a -> [a]
makePath Map a a
preds (Map a a
preds Map a a -> a -> a
forall k a. Ord k => Map k a -> k -> a
M.! a
v)

maxDisjointPaths :: (Mutable t, DirectedGraph t, Adjacency t, Integral a) => t a -> a -> a -> [[a]]
maxDisjointPaths :: t a -> a -> a -> [[a]]
maxDisjointPaths t a
d a
s a
t = [a
s a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a]
makePath a
v | a
v <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
outneighbors t a
d a
s, (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) a -> Map a a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a a
succs]
  where
    d' :: t a
d'  = (a -> t a -> t a) -> t a -> [a] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex (t a -> t a
forall (t :: * -> *) a. DirectedGraph t => t a -> t a
empty t a
d) ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
va -> a -> a
forall a. Num a => a -> a -> a
+a
1] | a
v <- t a -> [a]
forall (t :: * -> *) a. DirectedGraph t => t a -> [a]
vertices t a
d])
    d'' :: t a
d'' = ((a, a) -> t a -> t a) -> t a -> [(a, a)] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addArc t a
d' ([(a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) | a
v <- t a -> [a]
forall (t :: * -> *) a. DirectedGraph t => t a -> [a]
vertices t a
d] [(a, a)] -> [(a, a)] -> [(a, a)]
forall a. [a] -> [a] -> [a]
++ [(a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
va -> a -> a
forall a. Num a => a -> a -> a
+a
1, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
u) | (a
v,a
u) <- t a -> [(a, a)]
forall (t :: * -> *) a. DirectedGraph t => t a -> [(a, a)]
arcs t a
d])
    succs :: Map a a
succs = [(a, a)] -> Map a a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, a)] -> Map a a) -> [(a, a)] -> Map a a
forall a b. (a -> b) -> a -> b
$ Map (a, a) Bool -> [(a, a)]
forall k a. Map k a -> [k]
M.keys (Map (a, a) Bool -> [(a, a)]) -> Map (a, a) Bool -> [(a, a)]
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> Map (a, a) Bool -> Map (a, a) Bool
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Bool -> Bool
forall a. a -> a
id) (Map (a, a) Bool -> Map (a, a) Bool)
-> Map (a, a) Bool -> Map (a, a) Bool
forall a b. (a -> b) -> a -> b
$ t a -> a -> a -> Map (a, a) Bool
forall a (t :: * -> *).
(Ord a, Adjacency t, DirectedGraph t) =>
t a -> a -> a -> Map (a, a) Bool
maxFlow t a
d'' (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1) (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
t)
    makePath :: a -> [a]
makePath a
v
      | a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
t = [a
t]
      | Bool
otherwise = a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a]
makePath ((Map a a
succs Map a a -> a -> a
forall k a. Ord k => Map k a -> k -> a
M.! (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2)

minCut :: (Mutable t, DirectedGraph t, Adjacency t, Eq a) => t a -> a -> a -> [a]
minCut :: t a -> a -> a -> [a]
minCut t a
d a
s a
t = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Map Int a
iToV Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.!) ([Int] -> [a]) -> [Int] -> [a]
forall a b. (a -> b) -> a -> b
$ t Int -> Int -> Int -> [Int]
forall (t :: * -> *) a.
(Mutable t, DirectedGraph t, Adjacency t, Integral a) =>
t a -> a -> a -> [a]
minCutI t Int
di Int
si Int
ti
  where
    (t Int
di, [(Int, a)]
itova) = t a -> (t Int, [(Int, a)])
forall (t :: * -> *) a.
DirectedGraph t =>
t a -> (t Int, [(Int, a)])
linearizeVertices t a
d
    iToV :: Map Int a
iToV = [(Int, a)] -> Map Int a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int, a)]
itova
    Just Int
si = ((Int, a) -> Int) -> Maybe (Int, a) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, a) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, a) -> Maybe Int) -> Maybe (Int, a) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Bool) -> [(Int, a)] -> Maybe (Int, a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
s) (a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> a
forall a b. (a, b) -> b
snd) [(Int, a)]
itova
    Just Int
ti = ((Int, a) -> Int) -> Maybe (Int, a) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, a) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, a) -> Maybe Int) -> Maybe (Int, a) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Bool) -> [(Int, a)] -> Maybe (Int, a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
t) (a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> a
forall a b. (a, b) -> b
snd) [(Int, a)]
itova

minCutI :: (Mutable t, DirectedGraph t, Adjacency t, Integral a) => t a -> a -> a -> [a]
minCutI :: t a -> a -> a -> [a]
minCutI t a
d a
s a
t = [a
u a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2 | a
v <- Set a -> [a]
forall a. Set a -> [a]
S.toList (Set a -> [a]) -> Set a -> [a]
forall a b. (a -> b) -> a -> b
$ Set a
reach, a
u <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
outneighbors t a
d'' a
v, Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a
u a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
reach]
  where
    d' :: t a
d'  = (a -> t a -> t a) -> t a -> [a] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex (t a -> t a
forall (t :: * -> *) a. DirectedGraph t => t a -> t a
empty t a
d) ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
va -> a -> a
forall a. Num a => a -> a -> a
+a
1] | a
v <- t a -> [a]
forall (t :: * -> *) a. DirectedGraph t => t a -> [a]
vertices t a
d])
    d'' :: t a
d'' = ((a, a) -> t a -> t a) -> t a -> [(a, a)] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addArc t a
d' ([(a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
v a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) | a
v <- t a -> [a]
forall (t :: * -> *) a. DirectedGraph t => t a -> [a]
vertices t a
d] [(a, a)] -> [(a, a)] -> [(a, a)]
forall a. [a] -> [a] -> [a]
++ [(a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
va -> a -> a
forall a. Num a => a -> a -> a
+a
1, a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
u) | (a
v,a
u) <- t a -> [(a, a)]
forall (t :: * -> *) a. DirectedGraph t => t a -> [(a, a)]
arcs t a
d])
    flow :: Map (a, a) Bool
flow = (Bool -> Bool) -> Map (a, a) Bool -> Map (a, a) Bool
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Bool -> Bool
forall a. a -> a
id) (Map (a, a) Bool -> Map (a, a) Bool)
-> Map (a, a) Bool -> Map (a, a) Bool
forall a b. (a -> b) -> a -> b
$ t a -> a -> a -> Map (a, a) Bool
forall a (t :: * -> *).
(Ord a, Adjacency t, DirectedGraph t) =>
t a -> a -> a -> Map (a, a) Bool
maxFlow t a
d'' (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1) (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
t)
    reach :: Set a
reach = Set a -> Set a -> Set a
bfs (a -> Set a
forall a. a -> Set a
S.singleton (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1)) (a -> Set a
forall a. a -> Set a
S.singleton (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1))
    bfs :: Set a -> Set a -> Set a
bfs Set a
active Set a
reached
      | Set a -> Bool
forall a. Set a -> Bool
S.null Set a
active = Set a
reached
      | Bool
otherwise = Set a -> Set a -> Set a
bfs Set a
new (Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union Set a
reached Set a
new)
        where
          new :: Set a
new = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ [ a
u
                           | a
v <- Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
active
                           , a
u <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
outneighbors t a
d'' a
v
                           , (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (a
v,a
u) (a, a) -> Map (a, a) Bool -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map (a, a) Bool
flow) Bool -> Bool -> Bool
&& (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a
u a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
reached)
                           ]
                           [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
                           [ a
u
                           | a
v <- Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
active
                           , a
u <- t a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
inneighbors t a
d'' a
v
                           , (a
u,a
v) (a, a) -> Map (a, a) Bool -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map (a, a) Bool
flow Bool -> Bool -> Bool
&& (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a
u a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
reached)
                           ]