- Add the notes Samir emailed me on 2008-07-06.

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

diff --git a/program/BellmanFord.hs b/program/BellmanFord.hs
new file mode 100644 (file)
index 0000000..506acd9
--- /dev/null
+++ b/program/BellmanFord.hs
@@ -0,0 +1,69 @@
+module BellmanFord {-(spTree)-} where
+import Data.Graph.Inductive.Graph
+import Data.Graph.Inductive.Internal.Queue
+import Data.Graph.Inductive.Internal.RootPath
+import Data.Array.Diff
+
+data NodeInfo b = NodeInfo {
+       path :: Maybe [LNode b],
+       changed :: Bool
+}
+data Graph gr => BFState gr a b = BFState {
+       theGraph :: gr a b,
+       nis :: DiffArray Node (NodeInfo b),
+       changedQ :: Queue Node
+}
+
+nisToLRTree nis = do
+       ni <- elems nis
+       case path ni of
+               Just lnl -> return (LP lnl)
+               Nothing -> fail "Node is unreachable"
+
+offerPath :: (Graph gr, Real b) => BFState gr a b -> [LNode b] -> BFState gr a b 
+offerPath bfs newPath@((dest, newDist): _) =
+       -- Is newPath the first path to dest, or better than the previous one?
+       let adoptPath =
+               case path (nis bfs ! dest) of
+                       Nothing -> True
+                       Just ((_, oldDist) : _) -> newDist < oldDist
+               in
+       if adoptPath then
+               bfs{
+                       -- Update NodeInfo with the new path.
+                       nis = nis bfs // [(dest, NodeInfo (Just newPath) True)],
+                       changedQ = if changed (nis bfs ! dest)
+                               then changedQ bfs -- Already on the queue; leave as is.
+                               else queuePut dest (changedQ bfs) -- Add to queue.
+                       }
+               else bfs -- Don't update anything.
+
+processEdge :: (Graph gr, Real b) => [LNode b] -> LEdge b -> BFState gr a b -> BFState gr a b
+processEdge srcPath@((_, srcDist) : _) (_, dest, edgeLen) bfs =
+       let newPath = (dest, srcDist + edgeLen) : srcPath in
+       offerPath bfs newPath
+
+search :: (Graph gr, Real b) => BFState gr a b -> LRTree b
+search bfs =
+       if queueEmpty (changedQ bfs) then
+               -- Finished.
+               nisToLRTree (nis bfs)
+       else
+               -- Process a changed node from the queue.
+               let (src, moreQ) = queueGet (changedQ bfs) in
+               let srcNI = nis bfs ! src in
+               -- Clear src's changed flag.
+               let bfs1 = bfs{nis = nis bfs // [(src, srcNI{changed = False})], changedQ = moreQ} in
+               let Just srcPath = path srcNI in
+               let outEdges = out (theGraph bfs) src in
+               let newBFS = foldr (processEdge srcPath) bfs1 outEdges in
+               search newBFS
+
+spTree :: (Graph gr, Real b) => Node -> gr a b -> LRTree b
+spTree source theGraph = let theNodes = range (nodeRange theGraph) in
+       let emptyBFS = BFState theGraph
+               (array (nodeRange theGraph) (map (\n -> (n, NodeInfo Nothing False)) theNodes))
+               mkQueue in
+       -- Start with a zero-length path to the source.
+       let initBFS = offerPath emptyBFS [(source, 0)] in
+       search initBFS