leetcode - 1530. Number of Good Leaf Nodes Pairs

Description

You are given the root of a binary tree and an integer distance. A pair of two different leaf nodes of a binary tree is said to be good if the length of the shortest path between them is less than or equal to distance.

Return the number of good leaf node pairs in the tree.

Example 1:

Input: root = [1,2,3,null,4], distance = 3
Output: 1
Explanation: The leaf nodes of the tree are 3 and 4 and the length of the shortest path between them is 3. This is the only good pair.

Example 2:

Input: root = [1,2,3,4,5,6,7], distance = 3
Output: 2
Explanation: The good pairs are [4,5] and [6,7] with shortest path = 2. The pair [4,6] is not good because the length of ther shortest path between them is 4.

Example 3:

Input: root = [7,1,4,6,null,5,3,null,null,null,null,null,2], distance = 3
Output: 1
Explanation: The only good pair is [2,5].

Constraints:

The number of nodes in the tree is in the range [1, 2^10].
1 <= Node.val <= 100
1 <= distance <= 10

Solution

post-order visit

A pair of leaf nodes can only be formed if a node has both left and right child, so this node would work like a “bridge” to connect the leaf nodes from its left child and its right child. If given a distance d, we have the number of leaf nodes from left child l, and the number of leaf nodes from the right child r, then for the current node, it would have l x r pairs of leaves that meet the requirement.

So for each node, we maintain a hashmap to store the distance and the number of leaf nodes it has with this distance. And for a node with both children, we use a double loop to get all the pairs that meet the requirement.

Note that because the distance is small, we could use the distance for the double loop to reduce the time complexity.

Time complexity: $o(n\ast distanc{e}^2)$
Space complexity: $o(n\ast distance)$

Code

post-order visit

# 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: Optional[TreeNode], distance: int) -> int:
        # {node: {distance: cnt}}
        dis_info = {}
        stack = [(root, 0)]
        res = 0
        while stack:
            node, status = stack.pop()
            if status == 0:
                stack.append((node, 1))
                if node.left:
                    stack.append((node.left, 0))
                if node.right:
                    stack.append((node.right, 0))
            else:
                if node.left and node.right:
                    dis_info[node] = {}
                    for dis in range(distance):
                        if dis in dis_info[node.left]:
                            dis_info[node][dis + 1] = dis_info[node].get(dis + 1, 0) + dis_info[node.left][dis]
                        if dis in dis_info[node.right]:
                            dis_info[node][dis + 1] = dis_info[node].get(dis + 1, 0) + dis_info[node.right][dis]
                    for dis_i in range(distance):
                        for dis_j in range(distance):
                            if dis_i in dis_info[node.left] and dis_j in dis_info[node.right] and dis_i + dis_j + 2 <= distance:
                                res += dis_info[node.left][dis_i] * dis_info[node.right][dis_j]
                elif node.left:
                    dis_info[node] = {}
                    for dis in range(distance):
                        if dis in dis_info[node.left]:
                            dis_info[node][dis + 1] = dis_info[node].get(dis + 1, 0) + dis_info[node.left][dis]
                elif node.right:
                    dis_info[node] = {}
                    for dis in range(distance):
                        if dis in dis_info[node.right]:
                            dis_info[node][dis + 1] = dis_info[node].get(dis + 1, 0) + dis_info[node.right][dis]
                else:
                    dis_info[node] = {0: 1}
        return res