863. All Nodes Distance K in Binary Tree

Problem Description

The problem presents a binary tree and asks us to find all nodes that are a certain distance k from a given target node. It's important to note that distance here refers to the number of edges we must traverse to reach one node from another. This problem doesn't ask to find nodes only in the subtree of the target node but from the entire tree, so we must consider not only the children of the target node but also its ancestors and siblings.

Intuition

To solve this problem, we need to find all nodes that are at the specified distance k from the target node, regardless of whether they're ancestors, descendants, or neither (siblings or cousins). The approach to finding this solution is divided into several parts:

1. Mapping Parents: We first need a way to traverse the tree not only downwards from the root to the leaves (as is normally the case with tree traversals) but also upwards, from any node to its parent. This is not usually possible in a binary tree because nodes don't have a reference to their parent. Therefore, we create a map that keeps track of each node's parent.

2. Performing the Search: Once we have the ability to move both up and down the tree, we perform a depth-first search (DFS) starting from the target node. As we explore the tree, we keep track of the current distance from the target. When this distance equals k, we add the current node's value to our answer.

3. Avoiding Revisits: To ensure we don't count any nodes twice or enter a loop, we keep a set of visited nodes. Whenever we visit a node, we add it to the set. If we encounter a node that's already in our set, we skip it.

With this approach, we can explore all possible paths - down to the children, up to the parent, and across to siblings and other relatives - to find all nodes that are exactly k distance away from the target.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Approach

The solution is implemented using a depth-first search (DFS) algorithm along with additional data structures to track parent nodes and visited nodes. Here's a step-by-step breakdown of how the code accomplishes the task:

1. Parent Mapping: The first function, parents, is a recursive function that populates a dictionary p where keys are nodes of the tree and values are their respective parent nodes. It goes through all the nodes in a pre-order traversal (parent, then left child, then right child) and assigns the current node's parent to it in the dictionary.

2. Depth-First Search (DFS): The second function, dfs, is the crucial part of the implementation where we perform a depth-first search, starting from the target node. The function takes a node and the distance k as arguments.

3. Visited Set: Inside the DFS function, we first ensure that the current node is not None and hasn't been visited yet, as indicated by checking the vis set. If we have seen the node already, we return early to avoid revisiting it.

4. Base Case for Distance: If the distance k is 0, it means we have found a node at the exact distance from the target we are looking for, so we append the value of the current node to the ans list, which will eventually be our output.

5. Recursive Case for Traversal: If k is not 0, we recursively call DFS on the left child, the right child, and the parent of the current node, each time with k - 1 to account for the decrease in distance as we move one step away from the target.

6. Execution: Finally, we initialize our data structures: p for parents, ans for answers, and vis for visited nodes. We then populate p using the parents function and execute dfs starting with the target node and the initial distance k.

7. Result: After DFS completes the traversal of the tree, ans will contain all the node values that are k distance away from the target node. At the end of the function, we simply return ans.

In summary, the algorithm traverses the entire binary tree to build a parent-child relationship map, then it uses DFS to search all possible paths from the target node to find nodes that are exactly k steps away while avoiding loops and repetitions.

Example Walkthrough

Let's walk through a small example to illustrate the solution approach described above. Suppose we have the following binary tree and we want to find all the nodes that are a distance k = 2 away from the target node with value 5.

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

Step 1: Parent Mapping First, we construct a map of parent pointers:

• The parent of node 2 is 1.
• The parent of node 3 is 1.
• The parent of node 4 is 2.
• The parent of node 5 is 2, and so on.

Step 2: Performing the Search We then perform a DFS starting from the target node 5. We are looking for nodes that are at a distance k = 2.

Step 3: Avoiding Revisits We start with an empty vis set to keep track of visited nodes.

Starting from node 5, the search goes as follows:

• Visit node 5 (distance k = 2), mark it as visited.
• Move to node 8 (distance k = 1), mark it as visited.
• Move to node 4 (distance k = 0), since the distance is 0 and node 4 hasn't been visited, add it to the answer.
• Similarly, move to node 9 (distance k = 1), mark it as visited.
• Since node 9 has no children, return to node 5 and move up to its parent node 2 (distance k = 1), mark it as visited.
• Since node 2 is visited for the first time with distance k = 1, move to node 4 (already visited) and node 3 (distance k = 0 via parent 1). Again, as the distance is 0, add node 3 to the answer.

Step 4: Summary of the Results At the end of the DFS, we find that the nodes at distance 2 from the target node 5 are nodes 3 and 4. Thus, the answer is [3, 4].

By following the steps in the provided solution approach, we have successfully found all nodes at a specific distance from a given target in a binary tree.

Solution Implementation

Time and Space Complexity

The provided Python code defines a function distanceK that finds all nodes at distance k from a target node in a binary tree.

Time Complexity:

To analyze the time complexity, let's break down the procedure into its main steps:

1. The function parents traverses the entire tree to build a dictionary p mapping each node to its parent. This is a depth-first search (DFS) with a single visit to each node, hence its complexity is O(N), where N is the number of nodes in the tree.
2. The function dfs is a recursive depth-first search. In the worst case, it can visit each node once when k equals the height of the tree. Additionally, it might visit each node's parent, resulting in O(2N) operations, which simplifies to O(N).

Given these steps, the overall time complexity is O(N) since both parents and dfs functions operate linearly with respect to the number of nodes in the tree.

Space Complexity:

Now, let's consider the space complexity:

1. The p dictionary holds a pair for each node in the tree, hence O(N) space.
2. The ans list could hold up to N nodes (in the case that the tree is a straight line and k equals N-1), so its space is O(N).
3. The call stack for the DFS function can go up to the height of the tree in the case of a skewed tree, which in the worst case is O(N).
4. The vis set contains nodes that are already visited, which in the worst case can be O(N) as well when every node is visited.

The space complexity is determined by adding the space required for the p dictionary, the ans list, the call stack, and the vis set, which all contribute linearly to O(N) space.

Overall, taking into account both the worst-case scenarios for the recursive calls stack and the space needed for auxiliary data structures, the space complexity is O(N).

Learn more about how to find time and space complexity quickly using problem constraints.