BellmanFord.hs

   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