| 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 | ) |