Scala函数式编程Tree结构按序号获取从左到右第n个叶子节点问题

你当前的实现存在3个核心问题：

• 类型不匹配：函数声明返回值为A，但Branch分支返回的是Branch类型，编译阶段就无法通过
• 死循环风险：Leaf分支中当n != 0时，重复传入当前叶子节点递归，永远不会跳出递归
• 逻辑错误：遍历分支节点时没有统计左子树的叶子总数，直接给左右子树传入相同的n值，无法正确定位目标叶子所在的子树

按照从左到右数叶子的需求，优先遍历左子树，基于子树叶子总数做分支判断即可：

1. 先将用户传入的1起始的位置n转换为0起始的索引，简化计算逻辑
2. 遇到Leaf节点时直接返回节点值（进入该分支说明已经匹配到目标位置）
3. 遇到Branch节点时，先调用已有的count函数获取左子树的叶子总数：
   • 如果目标索引小于左子树叶子数，说明目标在左子树，递归查找左子树
   • 否则目标在右子树，将索引减去左子树叶子数后，递归查找右子树

// 基础Tree结构，和《Functional Programming in Scala》定义一致
sealed trait Tree[+A]
case class Leaf[A](value: A) extends Tree[A]
case class Branch[A](left: Tree[A], right: Tree[A]) extends Tree[A]

// 你已实现的count函数，这里补充作为参考
def count[A](tree: Tree[A]): Int = tree match {
  case Leaf(_) => 1
  case Branch(l, r) => count(l) + count(r)
}

// 查找第n个叶子的实现，n为1起始的位置编号
def getNthLeaf[A](tree: Tree[A], n: Int): A = {
  // 参数合法性校验
  val total = count(tree)
  require(n >= 1 && n <= total, s"n必须介于1到$total之间")
  // 内部递归使用0起始索引
  def loop(subTree: Tree[A], idx: Int): A = subTree match {
    case Leaf(v) => v
    case Branch(l, r) =>
      val leftLeafCount = count(l)
      if (idx < leftLeafCount) loop(l, idx)
      else loop(r, idx - leftLeafCount)
  }
  loop(tree, n - 1)
}

如果树结构较大，上述实现会因为多次调用count重复遍历子树，可以用单次遍历同时维护计数和查找，性能更优：

def getNthLeafOptimized[A](tree: Tree[A], n: Int): A = {
  var remaining = n // 对外保持纯函数性，仅内部用可变变量维护剩余计数
  def loop(subTree: Tree[A]): Option[A] = subTree match {
    case Leaf(v) =>
      remaining -= 1
      if (remaining == 0) Some(v) else None
    case Branch(l, r) =>
      loop(l).orElse(loop(r))
  }
  loop(tree).getOrElse(throw new IllegalArgumentException(s"n超出叶子节点总数范围"))
}

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