863. All Nodes Distance K in Binary Tree

MediumTreeDepth-First SearchBreadth-First SearchBinary 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.

Flowchart Walkthrough

Using the given flowchart, let's analyze and deduce the algorithm to use for LeetCode 863, All Nodes Distance K in Binary Tree. We'll conduct a step-by-step walkthrough:

Is it a graph?

Yes: The binary tree can be treated as a graph where each node is a vertex and edges exist between parent and child nodes.

Is it a tree?

Yes: Since the problem explicitly mentions that it deals with a binary tree, it fulfills the definition of a tree (acyclic connected graph).

DFS

Conclusion: Since we've identified it as a tree and there is no further specification to handle acyclic graphs, topological sort, or connectivity issues on shortest paths, the Depth-First Search (DFS) pattern is suggested for traversal or operation on tree structures. For this specific problem, DFS is apt for exploring all nodes to find those that are at a specific distance, effectively managing tree traversal from a given node.

Thus, the flowchart leads us to utilize the Depth-First Search algorithm for solving this problem regarding distances in a binary tree.

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:

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

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

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

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
         /   \
        2     3
       / \   / \
      4   5 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

1# Definition for a binary tree node.
2class TreeNode:
3    def __init__(self, val=0, left=None, right=None):
4        self.val = val
5        self.left = left
6        self.right = right
7
8
9class Solution:
10    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
11        # Helper function to store parent pointers for each node
12        def map_parents(node, parent):
13            if node:
14                parent_map[node] = parent
15                map_parents(node.left, node)
16                map_parents(node.right, node)
17
18        # Helper function to perform Depth-First Search (DFS) and find nodes at distance k
19        def find_nodes_at_distance_k(node, remaining_distance):
20            if not node or node.val in visited:
21                return
22            visited.add(node.val)
23            if remaining_distance == 0:
24                result.append(node.val)
25            else:
26                # Visit children and parent
27                find_nodes_at_distance_k(node.left, remaining_distance - 1)
28                find_nodes_at_distance_k(node.right, remaining_distance - 1)
29                find_nodes_at_distance_k(parent_map[node], remaining_distance - 1)
30
31        # Initialize the parent map and result list
32        parent_map = {}
33        result = []
34        visited = set()
35
36        # Map parents for all nodes in the tree
37        map_parents(root, None)
38
39        # Start DFS from the target node
40        find_nodes_at_distance_k(target, k)
41
42        return result
43

1import java.util.*;
2
3// Definition for a binary tree node.
4class TreeNode {
5    int val;
6    TreeNode left;
7    TreeNode right;
8    TreeNode(int x) { val = x; }
9}
10
11class Solution {
12    private Map<TreeNode, TreeNode> parentMap;
13    private Set<Integer> visited;
14    private List<Integer> answer;
15
16    // This method finds all nodes at distance K from the target node.
17    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
18        parentMap = new HashMap<>(); // Used to store parent-child relationships.
19        visited = new HashSet<>(); // Tracks visited nodes to prevent cycles.
20        answer = new ArrayList<>(); // Stores the results.
21
22        // Associates each node with its parent.
23        mapParents(root, null);
24      
25        // Perform DFS to find nodes at distance K.
26        findNodesAtDistanceK(target, k);
27      
28        return answer;
29    }
30
31    // Maps each node to its parent recursively.
32    private void mapParents(TreeNode node, TreeNode parent) {
33        if (node == null) {
34            return;
35        }
36      
37        parentMap.put(node, parent);
38        mapParents(node.left, node); // Parent for left child is 'node'
39        mapParents(node.right, node); // Parent for right child is 'node'
40    }
41
42    // DFS to find and add all nodes that are K distance from the given node.
43    private void findNodesAtDistanceK(TreeNode node, int k) {
44        if (node == null || visited.contains(node.val)) {
45            // Base case: if the node is null or already visited, stop the search.
46            return;
47        }
48      
49        visited.add(node.val); // Mark this node as visited.
50      
51        if (k == 0) {
52            // If the distance is zero, the current node is K distance from the target.
53            answer.add(node.val);
54            return; // Node added to the answer, end this path here.
55        }
56      
57        // Continue searching in the left subtree.
58        findNodesAtDistanceK(node.left, k - 1);
59      
60        // Continue searching in the right subtree.
61        findNodesAtDistanceK(node.right, k - 1);
62      
63        // Also search the sub-tree above using the parent references.
64        findNodesAtDistanceK(parentMap.get(node), k - 1);
65    }
66}
67

1/**
2 * Definition for a binary tree node.
3 */
4struct TreeNode {
5    int val;
6    TreeNode *left;
7    TreeNode *right;
8    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
9};
10
11class Solution {
12public:
13    // Maps each node to its parent
14    unordered_map<TreeNode*, TreeNode*> parentMap;
15
16    // Keeps track of visited nodes
17    unordered_set<int> visited;
18
19    // Stores the results
20    vector<int> result;
21
22    // Function that returns all nodes at distance K from the target node
23    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
24        // Find and store all parents of nodes
25        storeParents(root, nullptr);
26      
27        // Perform search to find all nodes at distance k
28        depthFirstSearch(target, k);
29      
30        return result;
31    }
32
33    // Helper method to relate each node to its parent in the tree
34    void storeParents(TreeNode* node, TreeNode* parent) {
35        if (!node) return;
36
37        parentMap[node] = parent;
38      
39        storeParents(node->left, node);
40        storeParents(node->right, node);
41    }
42
43    // The DFS method that traverses the tree to find nodes at a distance k
44    void depthFirstSearch(TreeNode* node, int k) {
45        // If the node is null or already visited, stop the search
46        if (!node || visited.count(node->val)) return;
47
48        // Mark the current node as visited
49        visited.insert(node->val);
50
51        // If the desired distance is reached, add the node's value to the results
52        if (k == 0) {
53            result.push_back(node->val);
54            return;
55        }
56
57        // Continue the search in left and right subtree
58        depthFirstSearch(node->left, k - 1);
59        depthFirstSearch(node->right, k - 1);
60
61        // Also consider the parent node in the search
62        depthFirstSearch(parentMap[node], k - 1);
63    }
64};
65

1// Definition for a binary tree node.
2class TreeNode {
3    val: number;
4    left: TreeNode | null;
5    right: TreeNode | null;
6
7    constructor(val: number) {
8        this.val = val;
9        this.left = null;
10        this.right = null;
11    }
12}
13
14// Maps each node to its parent
15const parentMap: Map<TreeNode, TreeNode | null> = new Map();
16
17// Keeps track of visited nodes
18const visited: Set<number> = new Set();
19
20// Stores the results
21const result: number[] = [];
22
23// Function that returns all nodes at distance K from the target node
24function distanceK(root: TreeNode | null, target: TreeNode, k: number): number[] {
25    // Find and store all parents of nodes
26    storeParents(root, null);
27  
28    // Perform search to find all nodes at distance k
29    depthFirstSearch(target, k);
30  
31    return result;
32}
33
34// Helper function to relate each node to its parent in the tree
35function storeParents(node: TreeNode | null, parent: TreeNode | null) {
36    if (!node) return;
37
38    parentMap.set(node, parent);
39  
40    storeParents(node.left, node);
41    storeParents(node.right, node);
42}
43
44// The DFS method that traverses the tree to find nodes at a distance k
45function depthFirstSearch(node: TreeNode | null, k: number) {
46    // If the node is null or already visited, stop the search
47    if (!node || visited.has(node.val)) return;
48
49    // Mark the current node as visited
50    visited.add(node.val);
51
52    // If the desired distance is reached, add the node's value to the results
53    if (k === 0) {
54        result.push(node.val);
55        return;
56    }
57
58    // Continue the search in the left and right subtree
59    depthFirstSearch(node.left, k - 1);
60    depthFirstSearch(node.right, k - 1);
61
62    // Also consider the parent node in the search
63    const parentNode = parentMap.get(node) || null;
64    depthFirstSearch(parentNode, k - 1);
65}
66

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:

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

The p dictionary holds a pair for each node in the tree, hence O(N) space.
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).
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).
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.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

How would you design a stack which has a function min that returns the minimum element in the stack, in addition to push and pop? All push, pop, min should have running time O(1).

Add another stack to store the minimums

Add a heap to store the minimums

Add a variable to track the minimum

Add a queue to store the minumums

I don't know