1530. 好叶子节点对的数量

1530. 好叶子节点对的数量

原始题目链接：https://leetcode-cn.com/problems/number-of-good-leaf-nodes-pairs/

给你二叉树的根节点 root 和一个整数 distance 。
如果二叉树中两个 叶 节点之间的 最短路径长度 小于或者等于 distance ，那它们就可以构成一组 好叶子节点对 。
返回树中 好叶子节点对的数量 。

输入：root = [1,2,3,null,4], distance = 3
输出：1
解释：树的叶节点是 3 和 4 ，它们之间的最短路径的长度是 3 。这是唯一的好叶子节点对。

输入：root = [1,2,3,4,5,6,7], distance = 3
输出：2
解释：好叶子节点对为 [4,5] 和 [6,7] ，最短路径长度都是 2 。但是叶子节点对 [4,6] 不满足要求，因为它们之间的最短路径长度为 4 。

解题思路：

两个叶子节点的最短路径 = 其中一个叶子节点到最近公共祖先的距离 + 另外一个叶子节点到最近公共祖先的距离
定义一个函数dfs(root)，root节点为当前遍历到的节点，计算以root节点为父节点其左右子树中的所有叶子节点到root节点的距离，因为返回的是所有到叶子节点的距离，所以需要用数组来存储。

代码实现：

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def __init__(self):
        self.ans = 0
    def countPairs(self, root: TreeNode, distance: int) -> int:
        # 定义一个函数dfs(root)：root节点为当前遍历到的节点，计算以root节点为父节点其左右子树中的所有叶子节点到root节点的距离
        # 因为返回的是所有到叶子节点的距离，所以需要用数组来存储
        def dfs(root):
            # 如果树不存在
            if not root:
                return []
            # 如果只有一个根节点的树
            if not root.left and not root.right:
                return [0]
            # 父节点到各个叶子节点的距离就是 父节点到子节点的距离+子节点到叶子节点的距离
            # 所以加1
            left = [l + 1 for l in dfs(root.left)]
            right = [r + 1 for r in dfs(root.right)]
            
            # 筛选符合的条件，笛卡尔积的形式计算组合结果
            for l in left:
                for r in right:
                    if l + r <= distance:
                        self.ans += 1
            # 向上返回当前节点的所有叶子节点距当前节点的距离给当前节点的上层节点
            # 所有叶子节点就是包含了左右子树的叶子节点的距离，即两个数组相加的结果
            return left + right
        dfs(root)
        return self.ans

参考文献：
https://leetcode-solution-leetcode-pp.gitbook.io/leetcode-solution/thinkings/tree#qi-ge-ji-qiao