Leetcode 1530. Number of Good Leaf Nodes Pairs (python)

题目

解法：

将树转换成图，然后从每个leaf node开始进行BFS。虽然这种解法没有这么快，但是我觉得很标准。所有判断有没有出现过的node都用字典实现O(1)的查找时间复杂度
要注意两点：

• 不同的节点值可能会相同
• (nodeA,nodeB)和(nodeB,nodeA)算一个pair

# 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
#         self.parent = None
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        leafs = {}
        def find_leaf(node,parent):
            if not node:
                return
            node.parent = parent
            if not node.left and not node.right:
                leafs[node] = True
            find_leaf(node.left,node)
            find_leaf(node.right,node)
        find_leaf(root,None)
    
        ans = []
        for leaf in leafs.keys():
            q = collections.deque()
            q.append((leaf,0))
            # visited = set()
            # visited.add(leaf)
            visited = {}
            visited[leaf] = True
            while q:
                node,dis = q.popleft()
                if node!=leaf and node in leafs and dis<=distance:
                    ans.append((leaf,node))
                for neigh in [node.left,node.right,node.parent]:
                    if neigh and neigh not in visited:
                        q.append((neigh,dis+1))
                        visited[neigh] = True
        count = 0
        seen = set()
        for a,b in ans:
            if (a,b) not in seen and (b,a) not in seen:
                count += 1
                seen.add((a,b))
        return count