title
给你二叉树的根节点 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
solution
class Solution:
def countPairs(self, root: TreeNode, distance: int) -> int:
ans = 0
def dfs(node):
'''
return all leaives and their depth
:param node:
:return:
'''
if not node: return []
if not node.left and not node.right: return [0]
l = dfs(node.left)
r = dfs(node.right)
nonlocal ans
res = [k + 1 for k in l if k < distance-1]
res += [k + 1 for k in r if k < distance-1]
for i in l:
for j in r:
if i + j+2<= distance:
ans += 1
return res
dfs(root)
return ans
优化后的,O(n*dis2)
class Solution:
def countPairs(self, root: TreeNode, distance: int) -> int:
ans = 0
def dfs(node):
'''
return all leaives and their depth
:param node:
:return:
'''
depth = [0 for _ in range(distance + 1)]
if not node: return depth
if not node.left and not node.right:
depth[0] = 1
return depth
l = dfs(node.left)
r = dfs(node.right)
nonlocal ans
for i in range(distance):
depth[i + 1] += l[i]
depth[i + 1] += r[i]
if l[i] > 0:
for j in range(distance - i - 1):
if r[j] > 0:
ans += (l[i] * r[j])
return depth
dfs(root)
return ans