| 1 | module BellmanFord {-(spTree)-} where |
| 2 | import Data.Graph.Inductive.Graph |
| 3 | import Data.Graph.Inductive.Internal.Queue |
| 4 | import Data.Graph.Inductive.Internal.RootPath |
| 5 | import Data.Array.Diff |
| 6 | |
| 7 | data NodeInfo b = NodeInfo { |
| 8 | path :: Maybe [LNode b], |
| 9 | changed :: Bool |
| 10 | } |
| 11 | data Graph gr => BFState gr a b = BFState { |
| 12 | theGraph :: gr a b, |
| 13 | nis :: DiffArray Node (NodeInfo b), |
| 14 | changedQ :: Queue Node |
| 15 | } |
| 16 | |
| 17 | nisToLRTree nis = do |
| 18 | ni <- elems nis |
| 19 | case path ni of |
| 20 | Just lnl -> return (LP lnl) |
| 21 | Nothing -> fail "Node is unreachable" |
| 22 | |
| 23 | offerPath :: (Graph gr, Real b) => BFState gr a b -> [LNode b] -> BFState gr a b |
| 24 | offerPath bfs newPath@((dest, newDist): _) = |
| 25 | -- Is newPath the first path to dest, or better than the previous one? |
| 26 | let adoptPath = |
| 27 | case path (nis bfs ! dest) of |
| 28 | Nothing -> True |
| 29 | Just ((_, oldDist) : _) -> newDist < oldDist |
| 30 | in |
| 31 | if adoptPath then |
| 32 | bfs{ |
| 33 | -- Update NodeInfo with the new path. |
| 34 | nis = nis bfs // [(dest, NodeInfo (Just newPath) True)], |
| 35 | changedQ = if changed (nis bfs ! dest) |
| 36 | then changedQ bfs -- Already on the queue; leave as is. |
| 37 | else queuePut dest (changedQ bfs) -- Add to queue. |
| 38 | } |
| 39 | else bfs -- Don't update anything. |
| 40 | |
| 41 | processEdge :: (Graph gr, Real b) => [LNode b] -> LEdge b -> BFState gr a b -> BFState gr a b |
| 42 | processEdge srcPath@((_, srcDist) : _) (_, dest, edgeLen) bfs = |
| 43 | let newPath = (dest, srcDist + edgeLen) : srcPath in |
| 44 | offerPath bfs newPath |
| 45 | |
| 46 | search :: (Graph gr, Real b) => BFState gr a b -> LRTree b |
| 47 | search bfs = |
| 48 | if queueEmpty (changedQ bfs) then |
| 49 | -- Finished. |
| 50 | nisToLRTree (nis bfs) |
| 51 | else |
| 52 | -- Process a changed node from the queue. |
| 53 | let (src, moreQ) = queueGet (changedQ bfs) in |
| 54 | let srcNI = nis bfs ! src in |
| 55 | -- Clear src's changed flag. |
| 56 | let bfs1 = bfs{nis = nis bfs // [(src, srcNI{changed = False})], changedQ = moreQ} in |
| 57 | let Just srcPath = path srcNI in |
| 58 | let outEdges = out (theGraph bfs) src in |
| 59 | let newBFS = foldr (processEdge srcPath) bfs1 outEdges in |
| 60 | search newBFS |
| 61 | |
| 62 | spTree :: (Graph gr, Real b) => Node -> gr a b -> LRTree b |
| 63 | spTree source theGraph = let theNodes = range (nodeRange theGraph) in |
| 64 | let emptyBFS = BFState theGraph |
| 65 | (array (nodeRange theGraph) (map (\n -> (n, NodeInfo Nothing False)) theNodes)) |
| 66 | mkQueue in |
| 67 | -- Start with a zero-length path to the source. |
| 68 | let initBFS = offerPath emptyBFS [(source, 0)] in |
| 69 | search initBFS |