LC 0863 [M] All Nodes Distance K in Binary Tree - ALawliet/algorithms GitHub Wiki

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