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] 。
提示：
- 树中节点数目在范围
[1, 2^10]内。 1 <= Node.val <= 1001 <= distance <= 10
解题思路
这是一道经典的二叉树深度优先搜索问题。主要思路是利用DFS后序遍历，对于每个节点，统计其子树中叶子节点到该节点的距离分布。
核心思想：
- 对于每个节点，我们需要知道其左右子树中叶子节点到该节点的距离分布
- 对于当前节点，好叶子节点对可能有两种情况：
- 完全在左子树中（递归处理）
- 完全在右子树中（递归处理）
- 一个在左子树，一个在右子树（需要在当前节点计算）
算法步骤：
- 使用DFS后序遍历，对于每个节点返回该子树中叶子节点到当前节点的距离数组
- 对于叶子节点，返回距离为0的数组
- 对于内部节点，合并左右子树的距离数组，距离都加1
- 在合并过程中，检查左右子树的叶子节点是否能构成好叶子节点对
- 如果左子树中距离为i的叶子节点和右子树中距离为j的叶子节点满足i+j+2≤distance，则它们构成好叶子节点对
这种方法的优势是只需要一次DFS遍历，时间复杂度较优。
代码实现
class Solution {
public:
int result = 0;
vector<int> dfs(TreeNode* node, int distance) {
if (!node) return {};
// 如果是叶子节点，返回距离为0
if (!node->left && !node->right) {
return {1}; // 距离0处有1个叶子节点
}
vector<int> left = dfs(node->left, distance);
vector<int> right = dfs(node->right, distance);
// 计算跨越当前节点的好叶子节点对
for (int i = 0; i < left.size(); i++) {
for (int j = 0; j < right.size(); j++) {
if (i + j + 2 <= distance) {
result += left[i] * right[j];
}
}
}
// 返回当前节点的距离分布（距离都加1）
vector<int> curr(distance, 0);
for (int i = 0; i < left.size() && i + 1 < distance; i++) {
curr[i + 1] += left[i];
}
for (int i = 0; i < right.size() && i + 1 < distance; i++) {
curr[i + 1] += right[i];
}
return curr;
}
int countPairs(TreeNode* root, int distance) {
dfs(root, distance);
return result;
}
};
class Solution:
def countPairs(self, root: Optional[TreeNode], distance: int) -> int:
self.result = 0
def dfs(node):
if not node:
return []
# 如果是叶子节点，返回距离为0
if not node.left and not node.right:
return [1] + [0] * (distance - 1)
left = dfs(node.left)
right = dfs(node.right)
# 计算跨越当前节点的好叶子节点对
for i in range(distance):
for j in range(distance):
if i + j + 2 <= distance:
self.result += left[i] * right[j]
# 返回当前节点的距离分布（距离都加1）
curr = [0] * distance
for i in range(distance - 1):
curr[i + 1] = left[i] + right[i]
return curr
dfs(root)
return self.result
public class Solution {
private int result = 0;
private int[] Dfs(TreeNode node, int distance) {
if (node == null) return new int[distance];
// 如果是叶子节点，返回距离为0
if (node.left == null && node.right == null) {
int[] leaf = new int[distance];
leaf[0] = 1;
return leaf;
}
int[] left = Dfs(node.left, distance);
int[] right = Dfs(node.right, distance);
// 计算跨越当前节点的好叶子节点对
for (int i = 0; i < distance; i++) {
for (int j = 0; j < distance; j++) {
if (i + j + 2 <= distance) {
result += left[i] * right[j];
}
}
}
// 返回当前节点的距离分布（距离都加1）
int[] curr = new int[distance];
for (int i = 0; i < distance - 1; i++) {
curr[i + 1] = left[i] + right[i];
}
return curr;
}
public int CountPairs(TreeNode root, int distance) {
result = 0;
Dfs(root, distance);
return result;
}
}
var countPairs = function(root, distance) {
let result = 0;
function dfs(node) {
if (!node) return new Array(distance).fill(0);
// 如果是叶子节点，返回距离为0
if (!node.left && !node.right) {
const leaf = new Array(distance).fill(0);
leaf[0] = 1;
return leaf;
}
const left = dfs(node.left);
const right = dfs(node.right);
// 计算跨越当前节点的好叶子节点对
for (let i = 0; i < distance; i++) {
for (let j = 0; j < distance; j++) {
if (i + j + 2 <= distance) {
result += left[i] * right[j];
}
}
}
// 返回当前节点的距离分布（距离都加1）
const curr = new Array(distance).fill(0);
for (let i = 0; i < distance - 1; i++) {
curr[i + 1] = left[i] + right[i];
}
return curr;
}
dfs(root);
return result;
};
复杂度分析
| 复杂度类型 | 分析 |
|---|---|
| 时间复杂度 | O(N × distance²)，其中 N 是树中节点的数量。对于每个节点，我们需要遍历距离数组进行计算，距离数组长度为 distance |
| 空间复杂度 | O(H × distance)，其中 H 是树的高度。递归调用栈的深度为 H，每次调用需要 O(distance) 的数组空间 |