Problem statement

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

Solution

Let us reconstruct our tree to graph, that is each node now can have 3 neighbours: left, right and parent. Then traverse graph once again, using bfs to found all nodes with distance K.

Complexity

Time and space complexity is O(n).

Code

class Solution:
    def distanceK(self, root, target, K):
        graph = defaultdict(list)
        
        def dfs(node):
            for child in [node.left, node.right]:
                if not child: continue
                graph[node.val].append(child.val)
                graph[child.val].append(node.val)
                dfs(child)
                
        queue = deque([(0, target.val)])
        visited, ans = {target.val}, []
                
        dfs(root)
        
        while queue:
            depth, node = queue.popleft()
            if depth == K: ans.append(node)
            for neib in graph[node]:
                if neib not in visited:
                    queue.append((depth + 1, neib))
                    visited.add(neib)
                    
        return ans