好叶子节点对的数量

难度: Medium

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

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

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

示例 1：

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

示例 2：

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

示例 3：

输入：root = [7,1,4,6,null,5,3,null,null,null,null,null,2], distance = 3
输出：1
解释：唯一的好叶子节点对是 [2,5] 。

示例 4：

输入：root = [100], distance = 1
输出：0

示例 5：

输入：root = [1,1,1], distance = 2
输出：1

提示：

tree 的节点数在 [1, 2^10] 范围内。
每个节点的值都在 [1, 100] 之间。
1 <= distance <= 10

Submission

运行时间: 63 ms

内存: 16.4 MB

class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        self.ans = 0

        def dfs(node, level = 1):
            if not node:
                return []

            if not node.left and not node.right:
                # 遇到叶子
                return [level]

            a = dfs(node.left, level + 1)
            b = dfs(node.right, level + 1)
            for it in a:
                t = it - level  # 边长
                if t >= distance:
                    break
                for jt in b:
                    t1 = jt - level
                    if t + t1 > distance:
                        break
                    else:
                        self.ans += 1
            # 合并两个有序数组，可优化。
            a.extend(b)
            a.sort()
            return a

        dfs(root)
        return self.ans

Explain

该题解采用了深度优先搜索（DFS）来遍历二叉树。在每个节点，如果它是叶子节点，则返回包含其深度的列表。如果不是叶子节点，则递归地处理左右子节点。将返回的左右子节点的深度列表进行比较，对于左子树的每个深度，检查右子树中与之可以配对形成好叶子节点对的深度（即两者深度之和不超过distance）。在遍历过程中统计这些好叶子节点对的数量。最终，返回总的好叶子节点对数。

时间复杂度: O(n^2 log n)

空间复杂度: O(n)

class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        self.ans = 0

        def dfs(node, level = 1):
            if not node:
                return []

            if not node.left and not node.right:
                # 这是一个叶子节点，返回包含当前深度的列表
                return [level]

            # 对左右子节点进行深度优先搜索
            a = dfs(node.left, level + 1)
            b = dfs(node.right, level + 1)
            # 检查左右叶子节点的深度组合是否形成好叶子节点对
            for it in a:
                t = it - level  # 计算当前叶子节点到父节点的深度
                if t >= distance:
                    break
                for jt in b:
                    t1 = jt - level
                    if t + t1 > distance:
                        break
                    else:
                        self.ans += 1
            # 合并并排序左右子树的深度列表
            a.extend(b)
            a.sort()
            return a

        dfs(root)
        return self.ans

Explore

在DFS函数中返回叶子节点的深度列表而不是直接计算好叶子节点对数的主要原因是为了将问题分解成可管理的部分，从而简化问题的复杂度。这种方法允许在每个节点处仅考虑其子树的叶子节点深度信息，而不是整个树的结构。这样，每个节点作为中介可以独立地评估与其左右子节点相关的所有可能的好叶子节点对。如果直接在每个节点下计算好叶子节点对数，将需要更复杂的状态管理和可能导致重复计数问题，因为相同的叶子对可能在树的多个级别被计算多次。通过返回深度列表，可以有效地在每个节点的上下文中只计算一次，且确保每对叶子只被统计一次。

题解中的break语句用于优化算法性能，避免不必要的比较。在内层循环中，当发现叶子节点对的总深度超过允许的最大distance时，break语句会终止当前的比较，因为右子树的深度列表是排序的，一旦当前组合深度超过distance，后续的组合也一定超过，因此没有继续检查的必要。在外层循环中，break语句的逻辑应是在深度已经达到或超过distance时停止循环，但从代码实现来看，外层循环的break语句实际上并未起到应有的效果，因为它是基于从父节点的深度计算，而不是从根节点。实际上，这个break可能不会在所有情况下正确地终止循环，特别是当左子树的叶子深度较小，但右子树深度较大且未排序的情况下。