{-# LANGUAGE ScopedTypeVariables #-}
module Algorithms.Graph.BFS
( bfs
, bfs'
) where
import Control.Monad.ST.Strict
import qualified Data.Foldable as F
import Data.Sequence (Seq(..))
import qualified Data.Sequence as Seq
import Data.Tree
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed.Mutable as UMV
import Witherable
bfs :: Foldable f => Int -> V.Vector (v, f Int) -> Tree v
bfs :: Int -> Vector (v, f Int) -> Tree v
bfs Int
s Vector (v, f Int)
gr = (Int -> v) -> Tree Int -> Tree v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((v, f Int) -> v
forall a b. (a, b) -> a
fst ((v, f Int) -> v) -> (Int -> (v, f Int)) -> Int -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector (v, f Int)
gr Vector (v, f Int) -> Int -> (v, f Int)
forall a. Vector a -> Int -> a
V.!)) (Tree Int -> Tree v)
-> (Vector (v, f Int) -> Tree Int) -> Vector (v, f Int) -> Tree v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector (f Int) -> Tree Int
forall (f :: * -> *).
Foldable f =>
Int -> Vector (f Int) -> Tree Int
bfs' Int
s (Vector (f Int) -> Tree Int)
-> (Vector (v, f Int) -> Vector (f Int))
-> Vector (v, f Int)
-> Tree Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((v, f Int) -> f Int) -> Vector (v, f Int) -> Vector (f Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v, f Int) -> f Int
forall a b. (a, b) -> b
snd (Vector (v, f Int) -> Tree v) -> Vector (v, f Int) -> Tree v
forall a b. (a -> b) -> a -> b
$ Vector (v, f Int)
gr
bfs' :: Foldable f => Int -> V.Vector (f Int) -> Tree Int
bfs' :: Int -> Vector (f Int) -> Tree Int
bfs' Int
s Vector (f Int)
gr = Int -> Vector [Int] -> Tree Int
extract Int
s (Vector [Int] -> Tree Int) -> Vector [Int] -> Tree Int
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MVector s [Int])) -> Vector [Int]
forall a. (forall s. ST s (MVector s a)) -> Vector a
V.create
((forall s. ST s (MVector s [Int])) -> Vector [Int])
-> (forall s. ST s (MVector s [Int])) -> Vector [Int]
forall a b. (a -> b) -> a -> b
$ do MVector s Bool
st <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UMV.replicate Int
n Bool
False
MVector s [Int]
out <- Int -> ST s (MVector (PrimState (ST s)) [Int])
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new Int
n
MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
forall s. MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
go0 MVector s Bool
st MVector s [Int]
out (Int
s Int -> Seq Int -> Seq Int
forall a. a -> Seq a -> Seq a
:<| Seq Int
forall a. Monoid a => a
mempty)
MVector s [Int] -> ST s (MVector s [Int])
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s [Int]
out
where
n :: Int
n = Vector (f Int) -> Int
forall a. Vector a -> Int
V.length Vector (f Int)
gr
go0 :: forall s. UMV.MVector s Bool -> MV.MVector s [Int]
-> Seq Int -> ST s ()
go0 :: MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
go0 MVector s Bool
st MVector s [Int]
out = Seq Int -> ST s ()
go
where
visit :: Int -> ST s (Maybe Int)
visit Int
i = do Bool
b <- MVector (PrimState (ST s)) Bool -> Int -> ST s Bool
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UMV.read MVector s Bool
MVector (PrimState (ST s)) Bool
st Int
i
MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UMV.write MVector s Bool
MVector (PrimState (ST s)) Bool
st Int
i Bool
True
Maybe Int -> ST s (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> ST s (Maybe Int)) -> Maybe Int -> ST s (Maybe Int)
forall a b. (a -> b) -> a -> b
$ if Bool
b then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i
go :: Seq Int -> ST s ()
go :: Seq Int -> ST s ()
go = \case
Seq Int
Empty -> () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
(Int
u:<|Seq Int
queue) -> do [Int]
ns <- (Int -> ST s (Maybe Int)) -> [Int] -> ST s [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Witherable t, Applicative f) =>
(a -> f (Maybe b)) -> t a -> f (t b)
wither Int -> ST s (Maybe Int)
visit ([Int] -> ST s [Int]) -> (f Int -> [Int]) -> f Int -> ST s [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Int -> [Int]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (f Int -> ST s [Int]) -> f Int -> ST s [Int]
forall a b. (a -> b) -> a -> b
$ Vector (f Int)
gr Vector (f Int) -> Int -> f Int
forall a. Vector a -> Int -> a
V.! Int
u
MVector (PrimState (ST s)) [Int] -> Int -> [Int] -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector s [Int]
MVector (PrimState (ST s)) [Int]
out Int
u [Int]
ns
Seq Int -> ST s ()
go (Seq Int
queue Seq Int -> Seq Int -> Seq Int
forall a. Semigroup a => a -> a -> a
<> [Int] -> Seq Int
forall a. [a] -> Seq a
Seq.fromList [Int]
ns)
extract :: Int -> V.Vector [Int] -> Tree Int
Int
s Vector [Int]
v = Int -> Tree Int
go Int
s
where
go :: Int -> Tree Int
go Int
i = Int -> Forest Int -> Tree Int
forall a. a -> Forest a -> Tree a
Node Int
i ((Int -> Tree Int) -> [Int] -> Forest Int
forall a b. (a -> b) -> [a] -> [b]
map Int -> Tree Int
go ([Int] -> Forest Int) -> [Int] -> Forest Int
forall a b. (a -> b) -> a -> b
$ Vector [Int]
v Vector [Int] -> Int -> [Int]
forall a. Vector a -> Int -> a
V.! Int
i)