All Nodes Distance K in Binary Tree

Problem Id: 863 Difficulty: Medium Tag: Tree Tag: Depth-first Search Tag: Breadth-first Search


Intuition

Convert the tree to a undirected graph and then do a BFS to get nodes whose distance to the target node is K.

Solution


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
from collections import defaultdict, deque


class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
        adj = defaultdict(list)
        self._dfs_get_adj(root, adj)
        queue = deque([[target.val, 0]])
        ans = []
        visited = set()
        while queue:
            node, distance = queue.popleft()
            if node in visited:
                continue
            visited.add(node)
            if distance == K:
                ans.append(node)
            else:
                for neighbour in adj[node]:
                    queue.append([neighbour, distance + 1])
        return ans

    def _dfs_get_adj(self, node, adj):
        if node.left:
            p, q = node.val, node.left.val
            adj[p].append(q)
            adj[q].append(p)
            self._dfs_get_adj(node.left, adj)
        if node.right:
            p, q = node.val, node.right.val
            adj[p].append(q)
            adj[q].append(p)
            self._dfs_get_adj(node.right, adj)