李洪强iOS经典面试题35-按层遍历二叉树的节点
问题
给你一棵二叉树,请按层输出其的节点值,即:按从上到下,从左到右的顺序。
例如,如果给你如下一棵二叉树:
3 / \ 9 20 / \ 15 7
输出结果应该是:
[ [3], [9,20], [15,7]]
代码模版
本题的 Swift 代码模版如下:
private 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 levelOrder(_ root: TreeNode?) -> [[Int]] { }}
解答
本题出自 LeetCode 第 102 题,是一个典型的有关遍历的题目。为了按层遍历,我们需要使用「队列」,来将每一层的节点先保存下来,然后再依次处理。
因为我们不但需要按层来遍历,还需要按层来输出结果,所以我在代码中使用了两个队列,分别名为 level
和 nextLevel
,用于保存不同层的节点。
最终,整个算法逻辑是:
- 判断输入参数是否是为空。
- 将根节点加入到队列
level
中。 - 如果
level
不为空,则:3.1 将level
加入到结果ans
中。3.2 遍历level
的左子节点和右子节点,将其加入到nextLevel
中。3.3 将nextLevel
赋值给level
,重复第 3 步的判断。 - 将
ans
中的节点换成节点的值,返回结果。
因为我们是用 Swift 来实现代码,所以我使用了一些 Swift 语言的特性。例如:队列中我们保存的是节点的数据结构,但是最终输出的时候,我们需要输出的是值,在代码中,我使用了 Swift 的函数式的链式调用,将嵌套数组中的元素类型做了一次变换,如下所示:
let ans = result.map { $0.map { $0.val }}
另外,我们也使用了 Swift 特有的 guard
关键字,来处理参数的特殊情况。
完整的参考代码如下:
//// Binary Tree Level Order Traversal.swift//// Created by Tang Qiao.//import Foundationprivate 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 }}private class Solution { func levelOrder(_ root: TreeNode?) -> [[Int]] { guard let root = root else { return [] } var result = [[TreeNode]]() var level = [TreeNode]() level.append(root) while level.count != 0 { result.append(level) var nextLevel = [TreeNode]() for node in level { if let leftNode = node.left { nextLevel.append(leftNode) } if let rightNode = node.right { nextLevel.append(rightNode) } } level = nextLevel } let ans = result.map { $0.map { $0.val }} return ans }}
微信中排版代码非常不便,所以上述代码也可以从我的 Gist 中找到:https://gist.github.com/tangqiaoboy/a7d24ac5ca266d650aaa91615f39ffae
完成这道题的同学,可以试着练习一下 LeetCode 的第 107 题,看看能不能只改动一行代码,就把 107 题也解决掉。
欢迎大家把自己的代码也放到 gist 上回复给我,Objective-C 或 Swift 的都行。
玩得开心!