Problem Statement

Given the root node of a Binary Tree, a target node and a distance k, return a list of all nodes that are a distance k from the target.

Explanation

When I first saw this problem, I panicked! I knew that this some how had to be a Breadth First Search problem because BFS from any node gives you the shortest path to every other node that is connected to the tree. Since every node is represented as a TreeNode with only left and right pointers, I would need to know the parent too for me to do a BFS from the target node and that was the reason for my panic. In fact, I even started doing a dry run of a Floyd-Warshall algorithm as this algorithm gives the shortest path between all pairs of nodes in a graph in O(N^3) time.

Depth First Search to the rescue! I later realized that I can use DFS to generate a map of a node to it’s parent and at this point, I should stop looking at the tree as a tree but as a graph.

For example, given the tree:

       3
    /     \
   5       1
  / \     / \
 6   2   0   8
    / \
   7   4

For target = 5 and K = 2, we need to return [7, 4, 1]. Once we have performed DFS and generated the map, we should now start looking at this tree as a graph as such:

6 - 5 - 3 - 1 - 8
    |       |
7 - 2       0
    |
    4

Now, we can perform BFS from the target node and whenever we encounter a node which is K distance away, we can append it to our results array. We can even have an optimization to break the BFS loop when we encounter a node at distance greater than K as BFS guarantees that we will process all nodes at a distance d before going to distance d + 1.

The code for this problem is as follows:

extension TreeNode: Hashable, CustomStringConvertible {
  public var description: String {
    return "\(val)"
  }
  
  public var hashValue: Int {
    return val
  }
  
  public static func == (lhs: TreeNode, rhs: TreeNode) -> Bool {
    return lhs.val == rhs.val
  }
}

func distanceK(_ root: TreeNode?, _ target: TreeNode?, _ K: Int) -> [Int] {
  guard let root = root, let target = target else {
    return [Int]()
  }
  
  var parentMap = [TreeNode:TreeNode?]()
  var result = [Int]()
  
  func dfs(currentNode: TreeNode?, parentNode: TreeNode?) {
    guard let currentNode = currentNode else {
      return
    }
    
    parentMap[currentNode] = parentNode
    dfs(currentNode: currentNode.left, parentNode: currentNode)
    dfs(currentNode: currentNode.right, parentNode: currentNode)
  }
  
  dfs(currentNode: root, parentNode: nil)
  
  func bfs(startNode: TreeNode) {
    var q = [TreeNode]()
    var seenSet = [TreeNode]()
    
    q.append(startNode)
    seenSet.append(startNode)
    
    var distanceMap = [TreeNode:Int]()
    distanceMap[startNode] = 0
    
    if K == 0 {
      result.append(startNode.val)
      return
    }
    
    while !q.isEmpty {
      let node = q.removeFirst()
      if let left = node.left, !seenSet.contains(left) {
        q.append(left)
        seenSet.append(left)
        
        distanceMap[left] = distanceMap[node]! + 1
        if distanceMap[left]! == K {
          result.append(left.val)
        }
      }
      
      if let right = node.right, !seenSet.contains(right) {
        q.append(right)
        seenSet.append(right)
        
        distanceMap[right] = distanceMap[node]! + 1
        if distanceMap[right]! == K {
          result.append(right.val)
        }
      }
      
      if let optionalParent = parentMap[node], let parent = optionalParent, !seenSet.contains(parent) {
        q.append(parent)
        seenSet.append(parent)
        
        distanceMap[parent] = distanceMap[node]! + 1
        if distanceMap[parent]! == K {
          result.append(parent.val)
        }
      }
    }
  }
  
  bfs(startNode: target)
  return result
}
Time Complexity

The time complexity of this problem is O(N) for DFS and O(N) for BFS which gives us a total of O(N).

Space Complexity

The space complexity is O(lg N) for the DFS recursive stack and O(N) for the BFS queue which gives us a total of O(N).

Solution is hosted on Github.