All Nodes Distance K in Binary Tree

Intuition

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)