Daily bit(e) of C++ | Find all nodes of distance K in a binary tree
Daily bit(e) of C++ #226, Common interview problem: Find all nodes of distance K in a binary tree.
Today we will look at a common C++ interview problem: Find all nodes of distance K in a binary tree.
Given a binary tree containing unique integer values, return all nodes that are K distance from the given node N.
For example, in the above tree, three nodes are distance two from node nine: {2,3,7}.
Before you continue reading the solution, I encourage you to try to solve it yourself. Here is a Compiler Explorer link with a couple of test cases: https://compiler-explorer.com/z/8recYWq9r.
Solution
One option is to translate the binary tree, essentially a directed graph, into an undirected graph, in which we can easily find k-distance nodes by applying a breadth-first search.
However, we have a simpler option based on the following observation. Consider a node with one of its children. If the child doesn’t lie on the path between the node and our target, its distance to the target is simply one more than the distance of the parent node. If it does lie on the path between the node and the target, its distance is one less.
If we explore the graph using pre-order traversal, we will also have a second guarantee that a node is only on the path if it is also on the path between the root and the target node.
Using these observations, we can construct a two-pass solution.
First, we find our target and initialize the distances for all nodes on the path between the target and the tree's root.
In the second pass, if we have a precomputed value for a node, we know that it is on the path, which allows us to distinguish between the two situations. Also, when we encounter a node with the appropriate distance, we remember it.