A Haskell implementation of efficient Graph BFS algorithm.

bfs.hs

#!/usr/bin/env cabal
{- cabal:
  build-depends: array, base, containers
-}

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

buildUG :: (Vertex, Vertex) -> [Edge] -> Graph
buildUG bounds edges = buildG bounds $ edges ++ map (\(x, y) -> (y, x)) edges

bfs :: Graph -> Vertex -> (UArray Vertex Vertex, UArray Vertex Int)
bfs graph start = runST $ do
  -- Initialize
  queue <- newSTRef $ singleton start
  visited <- newArray (bounds graph) False :: ST s (STUArray s Vertex Bool)
  writeArray visited start True
  parents <- newGenArray (bounds graph) pure :: ST s (STUArray s Vertex Vertex)
  distances <- newArray (bounds graph) (-1) :: ST s (STUArray s Vertex Int)
  writeArray distances start 0

  -- Run the BFS loop
  bfsLoop graph queue visited parents distances

  -- Return the results
  (,) <$> freeze parents <*> freeze distances

bfsLoop :: Graph -> STRef s (Seq Vertex) -> STUArray s Vertex Bool -> STUArray s Vertex Vertex -> STUArray s Vertex Int -> ST s ()
bfsLoop graph queue visited parents distances = do
  q <- readSTRef queue
  case viewl q of
    EmptyL -> pure ()
    v :< vs -> do
      writeSTRef queue vs
      let neighbors = graph ! v
      mapM_ (visit v) neighbors
      bfsLoop graph queue visited parents distances
  where
    visit parent child = do
      isVisited <- readArray visited child
      unless isVisited $ do
        writeArray visited child True
        writeArray parents child parent
        distance <- (+ 1) <$> readArray distances parent
        writeArray distances child distance
        modifySTRef' queue (|> child)

main :: IO ()
main = do
  let graph = buildUG (1, 9) [(1, 2), (1, 3), (2, 4), (2, 5), (3, 6)]
  print $ bfs graph 1