Rewrite Bellman-Ford and min-cost flow, especially to stop the latter from crashing.

[match/match.git] / program / NaiveMinCostFlow.hs

diff --git a/program/NaiveMinCostFlow.hs b/program/NaiveMinCostFlow.hs
new file mode 100644 (file)
index 0000000..1347837
--- /dev/null
+++ b/program/NaiveMinCostFlow.hs
@@ -0,0 +1,100 @@
+module NaiveMinCostFlow (minCostFlow) where
+import BellmanFord
+import MonadStuff
+import Data.Array.IArray
+import Data.Array.ST
+import Control.Monad
+import Control.Monad.ST
+import Data.Graph.Inductive.Graph
+import Data.Graph.Inductive.Internal.RootPath
+import Data.List
+
+data MCFEdge i f c = MCFEdge {
+       edgeIdx   :: i,
+       edgeCap   :: f,
+       edgeCost  :: c,
+       edgeIsRev :: Bool
+}
+data MCFState s gr a b i f c = MCFState {
+       mcfGraph  :: gr a (MCFEdge i f c),
+       mcfSource :: Node,
+       mcfSink   :: Node,
+       mcfFlow   :: STArray s i f
+}
+
+edgeCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
+       MCFEdge i f c -> ST s f
+edgeCapLeft state (MCFEdge i cap _ isRev) = do
+       fwdFlow <- readArray (mcfFlow state) i
+       return (if isRev then fwdFlow else cap - fwdFlow)
+
+edgePush :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
+       MCFEdge i f c -> f -> ST s ()
+edgePush state (MCFEdge i _ _ isRev) nf = do
+       oldFwdFlow <- readArray (mcfFlow state) i
+       let newFwdFlow = if isRev then oldFwdFlow - nf else oldFwdFlow + nf
+       writeArray (mcfFlow state) i newFwdFlow
+
+pathCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
+       (MCFEdge i f c, BFPath (MCFEdge i f c) c) -> ST s f
+pathCapLeft state (lastEdge, BFPath _ _ mFrom) = do
+       lastCL <- edgeCapLeft state lastEdge
+       case mFrom of
+               Nothing -> return lastCL
+               Just from -> do
+                       fromCL <- pathCapLeft state from
+                       return (min lastCL fromCL)
+
+augment :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
+       f -> BFPath (MCFEdge i f c) c -> ST s ()
+augment state augAmt (BFPath _ _ mFrom) = case mFrom of
+       Nothing -> nop
+       Just (lastEdge, path1) -> do
+               edgePush state lastEdge augAmt
+               augment state augAmt path1
+
+doFlow :: forall s gr a b i f c. (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
+       ST s ()
+doFlow state = do
+       filteredEdges <- filterM (\(_, _, l) -> do
+                       ecl <- edgeCapLeft state l
+                       return (ecl /= 0)
+               ) (labEdges (mcfGraph state))
+       let filteredGraph = mkGraph (labNodes (mcfGraph state)) filteredEdges :: gr a (MCFEdge i f c)
+       -- Why won't we get a negative cycle?  The original graph from the
+       -- proposal matcher is acyclic (so no negative cycle), and if we
+       -- created a negative cycle, that would contradict the fact that we
+       -- always augment along the lowest-cost path.
+       let mAugPath = bellmanFord edgeCost (mcfSource state) filteredGraph
+               ! (mcfSink state)
+       case mAugPath of
+               Nothing -> nop -- Done.
+               -- source /= sink, so augPasth is not empty.
+               Just augPath@(BFPath _ _ (Just from)) -> do
+                       augAmt <- pathCapLeft state from
+                       augment state augAmt augPath
+                       doFlow state
+
+minCostFlow :: forall s gr a b i f c. (Graph gr, Ix i, Real f, Real c) =>
+       (i, i)       -> -- Range of edge indices
+       (b -> i)     -> -- Edge label -> unique edge index
+       (b -> f)     -> -- Edge label -> flow capacity
+       (b -> c)     -> -- Edge label -> cost per unit of flow
+       gr a b       -> -- Graph
+       (Node, Node) -> -- (source, sink)
+       Array i f       -- ! edge index -> flow value
+minCostFlow idxBounds edgeIdx edgeCap edgeCost theGraph (source, sink) = runSTArray (do
+               let ourFlipF isRev l =
+                       MCFEdge (edgeIdx l) (edgeCap l)
+                               (if isRev then -(edgeCost l) else edgeCost l)
+                               isRev
+               let graph2 = mkGraph (labNodes theGraph) (concatMap
+                       (\(n1, n2, l) -> [ -- Capacity of reverse edge is never used.
+                               (n1, n2, MCFEdge (edgeIdx l) (edgeCap l) (  edgeCost l ) False),
+                               (n2, n1, MCFEdge (edgeIdx l)  undefined  (-(edgeCost l)) True )
+                       ]) $ labEdges theGraph) :: gr a (MCFEdge i f c)
+               flow <- newArray idxBounds 0
+               let state = MCFState graph2 source sink flow
+               doFlow state
+               return (mcfFlow state)
+       )