Swift计算树层级平均值时如何跳过nil值？代码遇测试用例失败

我来帮你排查这个问题～你的测试用例[10,5,15,null,null,6,20]对应的树结构第三层有两个有效节点（6和20），但你的函数没有统计到这一层，核心原因应该是层级遍历过程中没有正确收集所有非nil的子节点，或者在计算逻辑中错误地跳过了存在有效节点的层级。

问题分析 #

你当前的输出是[10.0,10.0]，说明函数只处理了前两层。大概率是在广度优先搜索（BFS）的实现中，处理子节点时的逻辑有误：比如可能你只收集了父节点同时存在左右子节点的情况，或者在队列处理时没有遍历完当前层的所有节点，导致第三层的节点没有被加入队列进行统计。

修复方案 #

下面提供两种正确的实现方式，分别基于BFS（推荐，层级遍历更直观）和DFS递归，都能正确处理含nil节点的情况：

1. 基于BFS的实现

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

func averageOfLevels(_ root: TreeNode?) -> [Double] {
    guard let root = root else { return [] }
    
    var result = [Double]()
    var nodeQueue = [TreeNode]()
    nodeQueue.append(root)
    
    while !nodeQueue.isEmpty {
        let levelNodeCount = nodeQueue.count
        var levelSum = 0
        var validNodeCount = 0
        
        // 遍历当前层级的所有节点
        for _ in 0..<levelNodeCount {
            let currentNode = nodeQueue.removeFirst()
            levelSum += currentNode.val
            validNodeCount += 1
            
            // 只将非nil的子节点加入队列，用于下一层遍历
            if let leftChild = currentNode.left {
                nodeQueue.append(leftChild)
            }
            if let rightChild = currentNode.right {
                nodeQueue.append(rightChild)
            }
        }
        
        // 计算当前层级的平均值并加入结果
        if validNodeCount > 0 {
            result.append(Double(levelSum) / Double(validNodeCount))
        }
    }
    
    return result
}

2. 基于DFS递归的实现

func averageOfLevelsDFS(_ root: TreeNode?) -> [Double] {
    var levelSums = [Int]()
    var levelNodeCounts = [Int]()
    
    func traverse(_ node: TreeNode?, level: Int) {
        guard let node = node else { return }
        
        // 如果是当前层级的第一个节点，初始化sum和count
        if level >= levelSums.count {
            levelSums.append(node.val)
            levelNodeCounts.append(1)
        } else {
            // 累加当前层级的sum和count
            levelSums[level] += node.val
            levelNodeCounts[level] += 1
        }
        
        // 递归遍历左右子节点，层级+1
        traverse(node.left, level: level + 1)
        traverse(node.right, level: level + 1)
    }
    
    traverse(root, level: 0)
    // 计算每一层的平均值
    return zip(levelSums, levelNodeCounts).map { Double($0) / Double($1) }
}

测试验证 #

用你的测试用例构建树：

// 构建测试树：10 -> 5、15；15 ->6、20
let node6 = TreeNode(6)
let node20 = TreeNode(20)
let node15 = TreeNode(15, node6, node20)
let node5 = TreeNode(5)
let root = TreeNode(10, node5, node15)

print(averageOfLevels(root)) // 输出：[10.0, 10.0, 13.0]
print(averageOfLevelsDFS(root)) // 输出：[10.0, 10.0, 13.0]

这样就能得到正确的结果了，核心是确保每一层的所有有效节点都被统计，不管其父节点的其他子节点是否为nil。

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