program/NaiveMinCostFlow.hs

   1 module NaiveMinCostFlow (minCostFlow) where
   2 import IMinCostFlow
   3 import BellmanFord
   4 import MonadStuff
   5 import Data.Array.IArray
   6 import Data.Array.ST
   7 import Control.Monad
   8 import Control.Monad.ST
   9 import Data.Graph.Inductive.Graph
  10 import Data.Graph.Inductive.Internal.RootPath
  11 import Data.List
  12
  13 data MCFEdge i f c = MCFEdge {
  14         edgeIdx   :: i,
  15         edgeCap   :: f,
  16         edgeCost  :: c,
  17         edgeIsRev :: Bool
  18 }
  19 data MCFState s gr a b i f c = MCFState {
  20         mcfGraph  :: gr a (MCFEdge i f c),
  21         mcfSource :: Node,
  22         mcfSink   :: Node,
  23         mcfFlow   :: STArray s i f
  24 }
  25
  26 edgeCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
  27         MCFEdge i f c -> ST s f
  28 edgeCapLeft state (MCFEdge i cap _ isRev) = do
  29         fwdFlow <- readArray (mcfFlow state) i
  30         return (if isRev then fwdFlow else cap - fwdFlow)
  31
  32 edgePush :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
  33         MCFEdge i f c -> f -> ST s ()
  34 edgePush state (MCFEdge i _ _ isRev) nf = do
  35         oldFwdFlow <- readArray (mcfFlow state) i
  36         let newFwdFlow = if isRev then oldFwdFlow - nf else oldFwdFlow + nf
  37         writeArray (mcfFlow state) i newFwdFlow
  38
  39 pathCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
  40         (MCFEdge i f c, BFPath (MCFEdge i f c) c) -> ST s f
  41 pathCapLeft state (lastEdge, BFPath _ _ mFrom) = do
  42         lastCL <- edgeCapLeft state lastEdge
  43         case mFrom of
  44                 Nothing -> return lastCL
  45                 Just from -> do
  46                         fromCL <- pathCapLeft state from
  47                         return (min lastCL fromCL)
  48
  49 augment :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
  50         f -> BFPath (MCFEdge i f c) c -> ST s ()
  51 augment state augAmt (BFPath _ _ mFrom) = case mFrom of
  52         Nothing -> nop
  53         Just (lastEdge, path1) -> do
  54                 edgePush state lastEdge augAmt
  55                 augment state augAmt path1
  56
  57 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 ->
  58         ST s ()
  59 doFlow state = do
  60         filteredEdges <- filterM (\(_, _, l) -> do
  61                         ecl <- edgeCapLeft state l
  62                         return (ecl /= 0)
  63                 ) (labEdges (mcfGraph state))
  64         let filteredGraph = mkGraph (labNodes (mcfGraph state)) filteredEdges :: gr a (MCFEdge i f c)
  65         -- Why won't we get a negative cycle?  The original graph from the
  66         -- proposal matcher is acyclic (so no negative cycle), and if we
  67         -- created a negative cycle, that would contradict the fact that we
  68         -- always augment along the lowest-cost path.
  69         let mAugPath = bellmanFord edgeCost (mcfSource state) filteredGraph
  70                 ! (mcfSink state)
  71         case mAugPath of
  72                 Nothing -> nop -- Done.
  73                 -- source /= sink, so augPasth is not empty.
  74                 Just augPath@(BFPath _ _ (Just from)) -> do
  75                         augAmt <- pathCapLeft state from
  76                         augment state augAmt augPath
  77                         doFlow state
  78
  79 -- We need to put the type parameters in scope for the mkGraph call.
  80 minCostFlow :: forall gr a b i f c. (Graph gr, Ix i, Real f, Real c) => MinCostFlowImpl1 gr a b i f c
  81 minCostFlow idxBounds edgeIdx edgeCap edgeCost theGraph (source, sink) = runSTArray (do
  82                 let ourFlipF isRev l =
  83                         MCFEdge (edgeIdx l) (edgeCap l)
  84                                 (if isRev then -(edgeCost l) else edgeCost l)
  85                                 isRev
  86                 let graph2 = mkGraph (labNodes theGraph) (concatMap
  87                         (\(n1, n2, l) -> [ -- Capacity of reverse edge is never used.
  88                                 (n1, n2, MCFEdge (edgeIdx l) (edgeCap l) (  edgeCost l ) False),
  89                                 (n2, n1, MCFEdge (edgeIdx l)  undefined  (-(edgeCost l)) True )
  90                         ]) $ labEdges theGraph) :: gr a (MCFEdge i f c)
  91                 -- Initialize only the slots of the flow array corresponding to
  92                 -- existing edges to zero to catch buggy callers that query
  93                 -- nonexistent edges.
  94                 flow <- newArray idxBounds undefined
  95                 sequence $ map (\(_, _, l) -> writeArray flow (edgeIdx l) 0)
  96                         (labEdges theGraph)
  97                 let state = MCFState graph2 source sink flow
  98                 doFlow state
  99                 return (mcfFlow state)
 100         )