863. All Nodes Distance K in Binary Tree - cocoder39/coco39_LC GitHub Wiki

863. All Nodes Distance K in Binary Tree

1. convert tree to graph. for each root->child, convert it to root->child and child->root
2. start from target node, applying bfs/dfs to find the nodes with desired distance

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        graph = collections.defaultdict(list)
        self.buildGraph(graph, root)

        queue = collections.deque()
        visited = set()
        if target:
            queue.append((target, 0))
            visited.add(target)
        
        res = []
        while queue:
            cur, dis = queue.popleft()
            if dis == k:
                res.append(cur.val)
            if dis > k:
                break # level by level so this works
            for nex in graph[cur]:
                if nex not in visited:
                    queue.append((nex, dis+1))
                    visited.add(nex)
        return res
    
    def buildGraph(self, graph, root):
        if not root:
            return

        if root.left:
            graph[root].append(root.left)
            graph[root.left].append(root)
            self.buildGraph(graph, root.left)
        
        if root.right:
            graph[root].append(root.right)
            graph[root.right].append(root)
            self.buildGraph(graph, root.right)