-
Notifications
You must be signed in to change notification settings - Fork 0
LC 0863 [M] All Nodes Distance K in Binary Tree
Code with Senpai edited this page Jan 6, 2022
·
10 revisions
Trees have direction (parent / child relationships) and don't contain cycles. They fit with in the category of Directed Acyclic Graphs (or a DAG). So Trees are DAGs with the restriction that a child can only have one parent.
use DFS to convert the tree to an undirected graph, then do a graph (not tree) BFS
class Solution:
def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
res = []
# convert tree (directed acyclic graph) to undirected graph adj list
# to do undirected we need to know the parent
adj = defaultdict(list)
def dfs(node, parent):
# undirected
if node and parent: # really just have to check parent because there is always a node
adj[node].append(parent)
adj[parent].append(node)
if node.left:
dfs(node.left, node)
if node.right:
dfs(node.right, node)
dfs(root, None)
# now we have a graph, now we can use simply BFS to calculate K distance from node
Q = deque([(target, 0)])
visited = set([target])
while Q:
node, distance = Q.popleft()
if distance == K:
res.append(node.val)
# adjacency list traversal
for neighbor in adj[node]:
if neighbor not in visited:
Q.append((neighbor, distance + 1))
visited.add(neighbor)
return res
# 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):
conn = defaultdict(list)
def connect(node, parent):
# both parent and child are not empty
if node and parent:
# building an undirected graph representation, assign the
# child value for the parent as the key and vice versa
conn[node.val].append(parent.val)
conn[parent.val].append(node.val)
# in-order traversal
if node.left: connect(node.left, node)
if node.right: connect(node.right, node)
# the initial parent node of the root is None
connect(root, None)
# start the breadth-first search from the target, hence the starting level is 0
bfs = [target.val]
seen = set(bfs)
# all nodes at (k-1)th level must also be K steps away from the target node
for _ in range(K):
# expand the list comprehension to strip away the complexity
level = []
for q_node_val in bfs:
for connected_node_val in conn[q_node_val]:
if connected_node_val not in seen:
level.append(connected_node_val)
bfs = level
# add all the values in bfs into seen
seen |= set(bfs)
return bfs
class Solution2:
def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
ans = []
# Return distance from node to target if exists, else -1
# Vertex distance: the # of vertices on the path from node to target
def dfs(node):
if not node:
return -1
elif node is target:
subtree_add(node, 0)
return 1
else:
L = dfs(node.left)
R = dfs(node.right)
if L != -1:
if L == K: ans.append(node.val)
subtree_add(node.right, L + 1)
return L + 1
elif R != -1:
if R == K: ans.append(node.val)
subtree_add(node.left, R + 1)
return R + 1
else:
return -1
# Add all nodes 'K - dist' from the node to answer.
def subtree_add(node, dist):
if not node:
return
elif dist == K:
ans.append(node.val)
else:
subtree_add(node.left, dist + 1)
subtree_add(node.right, dist + 1)
dfs(root)
return ans
footer