1530. Number of Good Leaf Nodes Pairs

Loading...LeetCode

Brute Force

周赛时候看到这题，半天也就想到用 LCA 做 brute force 解，但一直TLE。不过 brute force 其实是能写的，需要利用一下题目的这一条说明：

1 <= distance <= 10

distance 有上限，意味着寻找 LCA 的时候就不需要一直往上找了，超过 distance 即可放弃，因此这个过程可以被看作O(1)，思路就可以概括为：

1. 遍历树，找到所有叶节点，用hash map记录parent关系（用来在后面步骤中向上追溯）

2. 遍历所有叶节点，记录所有距离它小于distance的祖先节点，如果有符合条件的pair，则LCA必然在这些祖先节点里

3. 在还没遍历到的叶节点中搜索，要求：两个节点与LCA之和不大于distance

DFS分治

Tree 问题考虑下分治，确认要从左右子树分别传递过来什么数据（怎么治）。

容易想到的一个是「当前node到各个叶节点的距离」，但叶节点本身我们并不关心，只关心数量，所以可以直接记录「到当前节点距离为n的叶节点数量」，且利用distance有限这个条件，可以直接用数组来记录。

题目要找的 good leaf node pairs 在这里有个特点是左右子树各自内部的pairs是被确定的，往上传递的过程中不会再改变，新增来的pairs只可能是左子树的leaf 和右子树leaf形成的pair。所以需要往上传递的是当前子树记录的合格pairs数量。

Brute Force

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        res, parents, leaves = 0, dict(), []
        
        def dfs(c, p):
            if not c: return
            if not c.left and not c.right: leaves.append(c)
            parents[c] = p
            dfs(c.left, c)
            dfs(c.right, c)
        
        dfs(root, None)
        
        def ifGood(possibleLCA, node):
            i, n = 0, node
            # 寻找距离之和不大于distance的LCA
            while i < distance and n:
                if n in possibleLCA and possibleLCA[n] + i <= distance:
                    return 1
                # 当前node不是的话，就往上移动一层
                i += 1
                n = parents[n]
            return 0

        for i in range(len(leaves)):
            k, l = 0, leaves[i]
            ancestors = collections.defaultdict(int)
            # 记录当前距离当前叶节点小于distance的所有祖先节点
            while k < distance and l:
                ancestors[l] = k
                l = parents[l]
                k += 1
            # 在其他叶子节点里，找符合条件的叶节点（两个节点与LCA距离之和不大于distance的）
            for j in range(i+1, len(leaves)):
                res += ifGood(ancestors, leaves[j])
                
        return res

DFS分治

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        
        def dfs(node):
            disCount, pairCount = [0 for _ in range(distance+1)], 0
            if not node: return disCount, 0
            if not node.left and not node.right:
                disCount[0] = 1
                return disCount, 0
            countL, pairL = dfs(node.left)
            countR, pairR = dfs(node.right)
            for disL, nL in enumerate(countL):
                for disR, nR in enumerate(countR):
                    if disL + disR + 2 <= distance:
                        pairCount += nL * nR
            for i in range(distance):
                disCount[i+1] = countL[i] + countR[i]
            return disCount, pairL + pairR + pairCount
                
        return dfs(root)[1]