【二叉树】好叶子节点对的数量

原创

haocold 2022-10-24 18:36:24 博主文章分类：算法 ©著作权

文章标签 swift 算法二叉树子节点子树 文章分类 虚拟化云计算

©著作权归作者所有：来自51CTO博客作者haocold的原创作品，请联系作者获取转载授权，否则将追究法律责任

0x00 题目

给你二叉树的根节点 root 和一个整数 distance

如果二叉树中两个 叶 节点之间的 最短路径长度
小于或者等于 distance
那它们就可以构成一组 好叶子节点对

返回树中好叶子节点对的数量

0x01 思路

节点 root 的所有叶子节点包括
左子树的所有叶子节点
右子树的所有叶子节点
分别从左右叶子节点中拿出一个 进行相加再判断
符合则统计加 1
然后就是递归操作子节点了

0x02 解法

语言：Swift

树节点：TreeNode

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 countPairs(_ root: TreeNode?, _ distance: Int) -> Int {
    guard let root = root else { return 0 }
    var count: Int = 0
    
    func dfs(_ root: TreeNode?, _ distance: Int) -> [Int] {
        guard let root = root else { return [] }
        if root.left == nil && root.right == nil { return [0] }
        
        var res: [Int] = []
        
        // 左子树的所有叶子节点
        var left = dfs(root.left, distance)
        for i in 0..<left.count {
            // 从下层返回，路径加 1
            left[i] += 1
            if left[i] > distance {
                continue;
            }
            res.append(left[i])
        }
        
        // 右子树的所有叶子节点
        var right = dfs(root.right, distance)
        for i in 0..<right.count {
            // 从下层返回，路径加 1
            right[i] += 1
            if right[i] > distance {
                continue;
            }
            res.append(right[i])
        }
        
        for i in 0..<left.count {
            for j in 0..<right.count {
                if left[i]+right[j] <= distance {
                    count += 1
                }
            }
        }
        return res
    }
    
    _ = dfs(root, distance)
    
    return count
}