We are given a binary tree (with root node root), a target node, and an integer value K.
Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.
Example 1:
Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
Output: [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.
Note that the inputs “root” and “target” are actually TreeNodes. The descriptions of the inputs above are just serializations of these objects.
Note:
- The given tree is non-empty.
- Each node in the tree has unique values 0 <= node.val <= 500.
- The target node is a node in the tree.
- 0 <= K <= 1000.
Solutions
1. 当成图来解-BFS
在构造图的时候，
# Time Complexity: O(n)
# Space Complexity: O(n)
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class Solution:
def distanceK(self, root, target, K):
"""
:type root: TreeNode
:type target: TreeNode
:type K: int
:rtype: List[int]
"""
self.graph = collections.defaultdict(list)
self.build_graph(None, root)
queue = collections.deque()
queue.append(target.val)
visited = set([target.val])
res = []
k = 0
while queue and k <= K:
size = len(queue)
for i in range(size):
node_val = queue.popleft()
if k == K:
res.append(node_val)
for child in self.graph[node_val]:
if child in visited:
continue
queue.append(child)
visited.add(child)
k += 1
return res
def build_graph(self, parent, child):
if parent and child:
self.graph[parent.val].append(child.val)
self.graph[child.val].append(parent.val)
if child.left:
self.build_graph(child, child.left)
if child.right:
self.build_graph(child, child.right)
# Runtime: 36 ms, faster than 97.96% of Python3 online submissions for All Nodes Distance K in Binary Tree.
# Memory Usage: 13.9 MB, less than 62.50% of Python3 online submissions for All Nodes Distance K in Binary Tree.
2. 递归方法
# Time Complexity: O(n)
# Space Complexity: O(n)
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class Solution:
def distanceK(self, root, target, K):
"""
:type root: TreeNode
:type target: TreeNode
:type K: int
:rtype: List[int]
"""
self.res = []
self.traverse(root, target, K)
return self.res
# Returns the distance from root to target.
# Returns -1 if target does not in the tree.
def traverse(self, root, target, K):
if not root:
return -1
if root == target:
self.collect(target, K)
return 0
l = self.traverse(root.left, target, K)
r = self.traverse(root.right, target, K)
if l >= 0:
if l == K-1:
self.res.append(root.val)
self.collect(root.right, K-l-2)
return l + 1
if r >= 0:
if r == K-1:
self.res.append(root.val)
self.collect(root.left, K-r-2)
return r + 1
return -1
# Collect nodes that are d steps from root.
def collect(self, root, d):
if not root or d < 0:
return
if d == 0:
self.res.append(root.val)
self.collect(root.left, d-1)
self.collect(root.right, d-1)
# Runtime: 44 ms, faster than 73.33% of Python3 online submissions for All Nodes Distance K in Binary Tree.
# Memory Usage: 13.9 MB, less than 62.50% of Python3 online submissions for All Nodes Distance K in Binary Tree.