{-# LANGUAGE ScopedTypeVariables #-}
{-|
== Breadth first search

Sources:

* https://en.wikipedia.org/wiki/Breadth-first_search

-}

module BFS
    ( bfs
    ) where

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed
import Data.Graph

import Misc ( modifyArray )

-- | BFS on a graph, starting from the given source vertices. One tree per source is returned, which
-- contains all the vertices reached from the source before they could be reached by another.
-- Note that this is unlike Data.Graph.dfs, which returns one Tree for each connected component.
-- O(n + m), for a graph with n vertices and m edges.
bfs :: Graph -> [Vertex] -> Forest Vertex
bfs :: Graph -> [Vertex] -> Forest Vertex
bfs Graph
g [Vertex]
vs = (Vertex -> Tree Vertex) -> [Vertex] -> Forest Vertex
forall a b. (a -> b) -> [a] -> [b]
map Vertex -> Tree Vertex
toTree [Vertex]
vs where
    bnds :: (Vertex, Vertex)
bnds = Graph -> (Vertex, Vertex)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds Graph
g
    g' :: Graph
g' = (forall s. ST s (STArray s Vertex [Vertex])) -> Graph
forall i e. (forall s. ST s (STArray s i e)) -> Array i e
runSTArray ((forall s. ST s (STArray s Vertex [Vertex])) -> Graph)
-> (forall s. ST s (STArray s Vertex [Vertex])) -> Graph
forall a b. (a -> b) -> a -> b
$ do
        STUArray s Vertex Bool
vis :: STUArray s Vertex Bool <- (Vertex, Vertex) -> Bool -> ST s (STUArray s Vertex Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex, Vertex)
bnds Bool
False
        STArray s Vertex [Vertex]
ch :: STArray s Vertex [Vertex] <- (Vertex, Vertex) -> [Vertex] -> ST s (STArray s Vertex [Vertex])
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex, Vertex)
bnds []
        let go :: [Vertex] -> ST s ()
go [Vertex]
q = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Vertex] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vertex]
q) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ ([Vertex] -> Vertex -> ST s [Vertex])
-> [Vertex] -> [Vertex] -> ST s [Vertex]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [Vertex] -> Vertex -> ST s [Vertex]
f [] [Vertex]
q ST s [Vertex] -> ([Vertex] -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Vertex] -> ST s ()
go
            f :: [Vertex] -> Vertex -> ST s [Vertex]
f [Vertex]
q Vertex
u = ([Vertex] -> Vertex -> ST s [Vertex])
-> [Vertex] -> [Vertex] -> ST s [Vertex]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [Vertex] -> Vertex -> ST s [Vertex]
f' [Vertex]
q (Graph
gGraph -> Vertex -> [Vertex]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u) where
                f' :: [Vertex] -> Vertex -> ST s [Vertex]
f' [Vertex]
q Vertex
v = STUArray s Vertex Bool -> Vertex -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Bool
vis Vertex
v ST s Bool -> (Bool -> ST s [Vertex]) -> ST s [Vertex]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
m -> if Bool
m then [Vertex] -> ST s [Vertex]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Vertex]
q else [Vertex] -> Vertex -> ST s [Vertex]
add [Vertex]
q Vertex
v
                add :: [Vertex] -> Vertex -> ST s [Vertex]
add [Vertex]
q Vertex
v = do
                    STUArray s Vertex Bool -> Vertex -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Bool
vis Vertex
v Bool
True
                    STArray s Vertex [Vertex]
-> Vertex -> ([Vertex] -> [Vertex]) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray STArray s Vertex [Vertex]
ch Vertex
u (Vertex
vVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:)
                    [Vertex] -> ST s [Vertex]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
vVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
q) :: ST s [Vertex]
        [Vertex] -> (Vertex -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Vertex]
vs ((Vertex -> ST s ()) -> ST s ()) -> (Vertex -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Vertex
v -> STUArray s Vertex Bool -> Vertex -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Bool
vis Vertex
v Bool
True
        [Vertex] -> ST s ()
go [Vertex]
vs
        STArray s Vertex [Vertex] -> ST s (STArray s Vertex [Vertex])
forall (f :: * -> *) a. Applicative f => a -> f a
pure STArray s Vertex [Vertex]
ch
    toTree :: Vertex -> Tree Vertex
toTree Vertex
u = Vertex -> Forest Vertex -> Tree Vertex
forall a. a -> Forest a -> Tree a
Node Vertex
u (Forest Vertex -> Tree Vertex) -> Forest Vertex -> Tree Vertex
forall a b. (a -> b) -> a -> b
$ (Vertex -> Tree Vertex) -> [Vertex] -> Forest Vertex
forall a b. (a -> b) -> [a] -> [b]
map Vertex -> Tree Vertex
toTree (Graph
g'Graph -> Vertex -> [Vertex]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u)