题目
解法:
将树转换成图,然后从每个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:returnnode.parent = parentif not node.left and not node.right:leafs[node] = Truefind_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] = Truewhile 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] = Truecount = 0seen = set()for a,b in ans:if (a,b) not in seen and (b,a) not in seen:count += 1seen.add((a,b))return count