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 }

The time complexity of this problem is `O(N)`

for DFS and `O(N)`

for BFS which gives us a total of `O(N)`

.

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)`

.

Problem taken from Leetcode.

Solution hosted on Github.