1530. 好叶子节点对的数量

2021/03/06 leetcode 共 1359 字,约 4 分钟

title

problem link

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

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

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

示例 1:

img

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

示例 2:

img

输入: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

Search

    Table of Contents