###### 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.