Swift 解法详解：LeetCode 366《寻找二叉树的叶子节点》

摘要

这道题乍一看有点像是普通的二叉树遍历，但它其实是让我们把二叉树“分层剥离”。你可以想象一棵树，它的叶子节点最先掉落，掉光之后剩下的新叶子继续掉落，直到整棵树光秃秃的。我们要做的就是模拟这个过程，把每一轮掉落的叶子节点收集起来。

描述

题目要求：

• 给定一棵二叉树，反复删除叶子节点。
• 每一轮，收集当前所有叶子节点的值。
• 输出一个二维数组，数组中的每一项表示某一轮被删除的叶子。

举个例子：

输入: 
        1
       / \
      2   3
     / \
    4   5

输出: [[4,5,3],[2],[1]]

解释过程：

• 第一轮叶子节点是 [4,5,3]。
• 第二轮叶子节点是 [2]。
• 第三轮叶子节点是 [1]。

就好像一棵树的叶子一层层被风吹掉，直到只剩下树干。

题解答案

直觉上，我们可以模拟“每次删掉叶子”的过程，但实现起来会很麻烦，要反复扫描整个树。

更高效的做法是：把问题转换成“高度归类”。

• 对于树中的一个节点，它最终掉落的那一轮，其实就是它的“高度”。
• 什么是高度？可以理解为“离最近叶子节点的距离”。
• 比如叶子本身高度为 0，它的父节点高度为 1，以此类推。
• 我们只需要遍历一次树，计算每个节点的高度，然后把它们归类到对应的桶里即可。

题解代码分析

下面是 Swift 的实现：

import Foundation

// 定义二叉树节点
public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init(_ val: Int) {
        self.val = val
        self.left = nil
        self.right = nil
    }
}

class Solution {
    func findLeaves(_ root: TreeNode?) -> [[Int]] {
        var result = [[Int]]()
        
        // 递归函数：返回节点的高度
        @discardableResult
        func getHeight(_ node: TreeNode?) -> Int {
            guard let node = node else { return -1 }
            
            let leftHeight = getHeight(node.left)
            let rightHeight = getHeight(node.right)
            let height = max(leftHeight, rightHeight) + 1
            
            // 确保结果数组足够大
            if result.count <= height {
                result.append([Int]())
            }
            result[height].append(node.val)
            
            return height
        }
        
        _ = getHeight(root)
        return result
    }
}

// MARK: - Demo 演示
let root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left?.left = TreeNode(4)
root.left?.right = TreeNode(5)

let solution = Solution()
print(solution.findLeaves(root)) // 输出 [[4, 5, 3], [2], [1]]

代码拆解

1. TreeNode 定义：标准的二叉树节点。

2. getHeight 递归函数：

   • 叶子节点返回高度 0。
   • 每往上一层，高度就 +1。
   • 最后通过高度，把节点值放进 result[height]。

3. 收集结果：

   • 按照高度分组，相当于模拟了一轮一轮的“叶子掉落”。
   • result[0] 就是第一批掉落的叶子，result[1] 是第二批，以此类推。

示例测试及结果

运行 Demo，结果是：

[[4, 5, 3], [2], [1]]

再补充一些额外测试：

// 单节点树
let single = TreeNode(10)
print(solution.findLeaves(single)) 
// [[10]]

// 完全二叉树
let root2 = TreeNode(1)
root2.left = TreeNode(2)
root2.right = TreeNode(3)
root2.left?.left = TreeNode(4)
root2.left?.right = TreeNode(5)
root2.right?.left = TreeNode(6)
root2.right?.right = TreeNode(7)
print(solution.findLeaves(root2)) 
// [[4, 5, 6, 7], [2, 3], [1]]

时间复杂度

• 每个节点只遍历一次，所以是 O(n)，其中 n 是节点数。

空间复杂度

• 递归栈深度最坏是树的高度，平均情况 O(log n)，最坏情况 O(n)（比如链表树）。
• 结果数组的大小也是 O(n)。

总结

这道题的关键点在于把“删除叶子”的动态过程，转化成“高度归类”的静态计算。

• 直观方法：反复删叶子，效率低。
• 优雅方法：计算高度，一次搞定。

这个技巧其实挺实用的，在现实里也有场景：

• 比如在公司裁员，最底层的员工（叶子节点）最先被裁，裁完一批后再轮到中层，最后才是老板。
• 或者树形结构的数据清理，比如清除文件系统时，往往要从最深的文件开始，一层层删除，避免空目录残留。