All Nodes Distance K in Binary Tree

 

Problem Description: Given a binary tree, a target node in the tree, and an integer k, return all the nodes at distance k from the target node. You can assume that the given target node exists in the tree.

Example:

python
# Example usage of the function
root = [3,5,1,6,2,0,8,null,null,7,4]
target = 5
k = 2
result = distanceK(root, target, k)
# The nodes at distance 2 from the target node 5 are [7, 4]

1. Clarify the problem: Before starting, let's ask some clarifying questions to ensure we fully understand the problem requirements:

  • Can the binary tree contain duplicate values?
  • Can the target node be null or not present in the tree?
  • Can the distance k be negative or zero?

2. Analyze the problem: Let's break down the problem to better understand its components:

  • Input: A binary tree represented by its root node, a target node in the tree, and an integer k.
  • Output: A list of nodes at distance k from the target node.
  • Constraints: The target node is guaranteed to exist in the tree.

3. Design an algorithm: To find all the nodes at distance k from the target node, we can follow these steps:

  1. Perform a depth-first search (DFS) to build a parent mapping for each node in the binary tree.
  2. Perform another DFS to find all the nodes at distance k from the target node.
    • Initialize an empty list result to store the nodes at distance k.
    • Start DFS from the target node:
      • If k is 0, add the current node to result and return.
      • If k is greater than 0, recursively explore the left and right subtrees, but exclude the parent node.
      • Decrement k by 1 for the recursive calls.
    • Return result.

4. Explain your approach: We will perform two depth-first searches (DFS) to solve this problem. In the first DFS, we will build a parent mapping for each node in the binary tree. In the second DFS, starting from the target node, we will explore the left and right subtrees while decrementing k until reaching the desired distance. We will exclude the parent node during exploration to avoid revisiting nodes. Finally, we will return the list of nodes at distance k.

5. Write clean and readable code:

Here's an implementation of the algorithm in Python:

python
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def buildParentMapping(root, parent_map):
    if root:
        if root.left:
            parent_map[root.left.val] = root
            buildParentMapping(root.left, parent_map)
        if root.right:
            parent_map[root.right.val] = root
            buildParentMapping(root.right, parent_map)


def distanceK(root, target, k):
    parent_map = {}
    buildParentMapping(root, parent_map)
    visited = set()
    result = []

    def dfs(node, dist):
        if not node:
            return
        if node in visited:
            return
        visited.add(node)
        if dist == k:
            result.append(node.val)
            return
        if node.left and node.left not in visited:
            dfs(node.left, dist + 1)
        if node.right and node.right not in visited:
            dfs(node.right, dist + 1)
        parent = parent_map.get(node.val)
        if parent and parent not in visited:
            dfs(parent, dist + 1)

    dfs(target, 0)
    return result

6. Test your code: Let's test the code with some example test cases:

python
# Example test case
root = TreeNode(3)
root.left = TreeNode(5)
root.right = TreeNode(1)
root.left.left = TreeNode(6)
root.left.right = TreeNode(2)
root.right.left = TreeNode(0)
root.right.right = TreeNode(8)
root.left.right.left = TreeNode(7)
root.left.right.right = TreeNode(4)
target = root.left
k = 2
result = distanceK(root, target, k)
# Expected output: [7, 4]

# Additional test cases
root = TreeNode(0)
target = root
k = 0
result = distanceK(root, target, k)
# Expected output: [0]

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
target = root
k = 2
result = distanceK(root, target, k)
# Expected output: []

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
target = root.left
k = 1
result = distanceK(root, target, k)
# Expected output: [3]

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.right = TreeNode(6)
target = root.left
k = 2
result = distanceK(root, target, k)
# Expected output: [1, 6]

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.right = TreeNode(6)
target = root.right.right
k = 1
result = distanceK(root, target, k)
# Expected output: [3]

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.right = TreeNode(6)
target = root
k = 3
result = distanceK(root, target, k)
# Expected output: []

7. Optimize if necessary: The implemented algorithm has a time complexity of O(n), where n is the number of nodes in the binary tree. Each node is visited exactly once during the two depth-first searches. This is already an optimal solution for this problem, so there is no further optimization required.

8. Handle error cases: The given code assumes that the input binary tree is valid and that the target node exists in the tree. It does not handle the cases where the root or target node is null. Error handling strategies can be discussed with the interviewer.

9. Discuss complexity analysis:

  • Time complexity: The algorithm has a time complexity of O(n), where n is the number of nodes in the binary tree. Each node is visited exactly once during the two depth-first searches.
  • Space complexity: The space complexity is O(n) since the algorithm uses additional space for the parent mapping and the visited set, both of which can contain up to n nodes in the worst case.
Next Post Previous Post