All Nodes Distance K in Binary Tree

Patrick Leaton
Published: July 23, 2021, 6:39 p.m.
Problem Description

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node.  The answer can be returned in any order.

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output: [7,4,1]

Explanation: 
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.



Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

 

Note:

  1. The given tree is non-empty.
  2. Each node in the tree has unique values 0 <= node.val <= 500.
  3. The target node is a node in the tree.
  4. 0 <= K <= 1000.

 

The description was taken from https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/.

Problem Solution

#O(N) Time, O(N) Space
class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
        if not root:
            return root
 
        adjacency, seen = {root.val:[]}, {}
        output = []
       
        def dfs(node:TreeNode, parent:TreeNode) -> None:
            if not node:
                return
            if node.val not in adjacency:
                adjacency[node.val] = []
            adjacency[parent.val].append(node.val)
            adjacency[node.val].append(parent.val)
            dfs(node.left, node)
            dfs(node.right, node)
       
        def bfs(target: int, level = 0) -> None:
            queue = collections.deque([(target, level)])
            while queue:
                node, node_level = queue.popleft()
                seen[node] = True
                if node_level == K:
                    output.append(node)
                for neighbor in adjacency[node]:
                    if neighbor not in seen:
                        seen[neighbor] = True
                        queue.append((neighbor, node_level + 1))      
        
        dfs(root.left, root)
        dfs(root.right, root)       
        bfs(target.val)
 
        return output

Problem Explanation

The goal is to retrieve all nodes on the kth level from the target.

Since we are doing a search of levels, we know that we should consider doing a Breadth-First Search starting from the target and keeping track of how many levels we have moved from the target.

However, the input tree is a directed graph, so we can only move downward and not upward.  To handle this, we can do an initial Depth-First Search to build an adjacency list, a map of neighbors for each node, so that we can also move toward parent nodes in our BFS.

First off, if the root is empty then we have no nodes distance k, so we are done.

        if not root:
            return root

 

Otherwise, we will initialize our adjacency list with the root's value, and an empty list of neighbors.

We will also initialize a seen dictionary to avoid cycles in the graph and an output array.

        adjacency, seen = {root.val:[]}, {}
        output = []

Next, we will make our recursive, dfs function that we will use to build the adjacency list.

Within each call, we will keep track of the parent node and the current node, so that we can create a bridge to the parent within our adjacency list.

        def dfs(node:TreeNode, parent:TreeNode) -> None:

 

If the current node is none, then we have reached the end of a valid path so we will return.

            if not node:
                return

 

If the node's value currently isn't in the adjacency list then we will put it there with an initial list of empty neighbors.

            if node.val not in adjacency:
                adjacency[node.val] = []

 

Next, we will create an adjacency relationship, a bridge, from the node to the parent and the parent to the node.

            adjacency[parent.val].append(node.val)
            adjacency[node.val].append(parent.val)

 

Afterward, we will continue the DFS on the left and right children and passing the current node as the parent within the function call.

            dfs(node.left, node)
            dfs(node.right, node)

Next, we will create our iterative, bfs function.

The function will take a starting node to start the BFS from and it will initialize the level counter at zero.

        def bfs(target: int, level = 0) -> None:

 

Let's initialize a queue and put a tuple inside of it containing the target node value and the starting level.

            queue = collections.deque([(target, level)])

 

While the queue still has nodes that we need to visit, we will continue to pop a node, visit it, visit its unseen neighbors, then push its neighbors on the queue while incrementing the level to repeat the process.

If the level counter signifies that the node is k levels away from the target, then we will append the node value to the output.

            while queue:
                node, node_level = queue.popleft()
                seen[node] = True
                if node_level == K:
                    output.append(node)
                for neighbor in adjacency[node]:
                    if neighbor not in seen:
                        seen[neighbor] = True
                        queue.append((neighbor, node_level + 1))      

Now that we have the dfs and bfs functions built, we'll make the initial DFS calls on the root's left and right children, with the root being the first parent whose value was initialized in the adjacency list previously.

        dfs(root.left, root)
        dfs(root.right, root)      

 

Then, we can run the BFS from the target value to get all nodes k levels away from the target.

        bfs(target.val)

 

Once we have done that, we can return the output containing all of those node values.

        return output