主体部分,基于Leetcode 1530题
# 获取当前节点左右子树中叶子节点到该点的距离
def DFS(root):
if not root:
return {
}
if not root.left and not root.right:
return {
root: 0}
left_leaf = DFS(root.left)
right_leaf = DFS(root.right)
for k, v in left_leaf.items():
left_leaf[k] = v + 1
for k, v in right_leaf.items():
right_leaf[k] = v + 1
# 组合这些叶子节点,判断是否存在合理情况
for kl, vl in left_leaf.items():
for kr, vr in right_leaf.items():
if vl + vr <= distance:
self.res += 1
# 合并左右子树叶子节点,向上返回
left_leaf.update(right_leaf)
return left_leaf