如何以Scala函数式风格重写源节点到各节点的距离计算函数？

问题描述 #

我编写了一个用于计算源节点到各节点距离的函数，节点为字符串类型，neighbours(node)函数可返回节点的邻居。该函数运行正常，但我希望它更符合Scala的风格，请问如何重写它？（或许可以用递归实现？）

当前实现代码如下：

def distanceMap(source: String): Map[String, Int] = {
    val distances = scala.collection.mutable.Map.empty[String, Int]
    var batch = List(source)
    var newBatch: List[String] = Nil
    val seen = scala.collection.mutable.Set(source)
    var distance = 1
    while (!batch.isEmpty) {
        newBatch = batch.flatMap(neighbours(_)).filterNot(seen(_))
        for (neighbour <- newBatch) {
            seen.add(neighbour)
            distances(neighbour) = distance
        }
        batch = newBatch
        distance += 1
    }
    distances.toMap
}

纯函数式递归实现（尾递归优化） #

Scala推崇纯函数、不可变数据结构，我们可以用尾递归替代while循环，完全避免可变变量，写出更符合Scala风格的代码：

import scala.annotation.tailrec

def distanceMap(source: String): Map[String, Int] = {
  // 尾递归辅助函数，维护已访问集合、当前批次节点、当前距离、累计结果Map
  @tailrec
  def bfs(visited: Set[String], currentBatch: List[String], currentDist: Int, acc: Map[String, Int]): Map[String, Int] = {
    if (currentBatch.isEmpty) acc
    else {
      // 获取当前批次所有未被访问的邻居
      val nextNodes = currentBatch.flatMap(neighbours).filterNot(visited)
      // 去重，避免同一节点重复处理
      val uniqueNextNodes = nextNodes.toSet.toList
      // 更新距离Map，新节点的距离为当前距离
      val newAcc = acc ++ uniqueNextNodes.map(_ -> currentDist)
      // 递归处理下一批次，更新已访问集合、批次节点、距离和结果
      bfs(visited ++ uniqueNextNodes, uniqueNextNodes, currentDist + 1, newAcc)
    }
  }

  // 初始调用：已访问集合包含源节点，当前批次为源节点，初始距离1，初始结果为空Map（与原函数行为一致，不含源节点自身距离）
  bfs(Set(source), List(source), 1, Map.empty)
}

关键说明 #

1. 尾递归优化：通过@tailrec注解标记辅助函数，编译器会将其转换为底层循环，避免栈溢出问题，性能与原while循环相当。
2. 不可变数据：使用不可变的Set和Map，每次递归都生成新的集合实例，完全符合函数式编程的无副作用要求。
3. 行为一致性：上述实现与原函数返回结果完全一致——不包含源节点到自身的距离。如果需要添加该距离，只需将初始结果Map改为Map(source -> 0)：

   bfs(Set(source), List(source), 1, Map(source -> 0))
   

4. 效率优化：用toSet.toList替代distinct做去重，处理大量节点时效率更高。

内容的提问来源于stack exchange，提问作者Weier