Functional Breadth First Search in Scala with the State Monad

Viewed 777

I'm trying to implement a functional Breadth First Search in Scala to compute the distances between a given node and all the other nodes in an unweighted graph. I've used a State Monad for this with the signature as :-

case class State[S,A](run:S => (A,S))

Other functions such as map, flatMap, sequence, modify etc etc are similar to what you'd find inside a standard State Monad.

Here's the code :-

case class Node(label: Int)

case class BfsState(q: Queue[Node], nodesList: List[Node], discovered: Set[Node], distanceFromSrc: Map[Node, Int]) {
  val isTerminated = q.isEmpty
}

case class Graph(adjList: Map[Node, List[Node]]) {
  def bfs(src: Node): (List[Node], Map[Node, Int]) = {
    val initialBfsState = BfsState(Queue(src), List(src), Set(src), Map(src -> 0))
    val output = bfsComp(initialBfsState)
    (output.nodesList,output.distanceFromSrc)
  }

  @tailrec
  private def bfsComp(currState:BfsState): BfsState = {
     if (currState.isTerminated) currState
     else bfsComp(searchNode.run(currState)._2)
  }

  private def searchNode: State[BfsState, Unit] = for {
    node <- State[BfsState, Node](s => {
      val (n, newQ) = s.q.dequeue
      (n, s.copy(q = newQ))
    })
    s <- get
    _ <- sequence(adjList(node).filter(!s.discovered(_)).map(n => {
      modify[BfsState](s => {
        s.copy(s.q.enqueue(n), n :: s.nodesList, s.discovered + n, s.distanceFromSrc + (n -> (s.distanceFromSrc(node) + 1)))
      })
    }))
  } yield ()
}   

Please can you advice on :-

  1. Should the State Transition on dequeue in the searchNode function be a member of BfsState itself?
  2. How do I make this code more performant/concise/readable?
1 Answers
Related