Tarjan's algorithm for Scala

TarjanRecursive.scala

package tarjan

import scala.collection.mutable

// See python version: https://github.com/bwesterb/py-tarjan
/**
  * Tarjan's algorithm is an algorithm in graph theory for finding the strongly connected components of a directed graph. It runs in linear time,
  * matching the time bound for alternative methods including Kosaraju's algorithm and the path-based strong component algorithm.
  */
class TarjanRecursive {

  /**
    * The algorithm takes a directed graph as input, a
    * nd produces a partition of the graph's vertices into the graph's strongly connected components.
    * @param g the graph
    * @return the strongly connected components of g
    */
  def apply(g: Map[Int, List[Int]]): mutable.Buffer[mutable.Buffer[Int]] = {
    val s = mutable.Buffer.empty[Int]
    val s_set = mutable.Set.empty[Int]
    val index = mutable.Map.empty[Int, Int]
    val lowlink = mutable.Map.empty[Int, Int]
    val ret = mutable.Buffer.empty[mutable.Buffer[Int]]
    def visit(v: Int): Unit = {
      index(v) = index.size
      lowlink(v) = index(v)
      s += v
      s_set += v
      for (w <- g(v)) {
        if (!index.contains(w)) {
          visit(w)
          lowlink(v) = math.min(lowlink(w), lowlink(v))
        } else if (s_set(w)) {
          lowlink(v) = math.min(lowlink(v), index(w))
        }
      }
      if (lowlink(v) == index(v)) {
        val scc = mutable.Buffer.empty[Int]
        var w = -1
        while(v != w) {
          w = s.remove(s.size - 1)
          scc += w
          s_set -= w
        }
        ret += scc
      }
    }
    for (v <- g.keys) if (!index.contains(v)) visit(v)
    ret
  }

}