Solving "All Nodes Distance K in Binary Tree" Leet code Question

Problem Statement

Given a binary tree, you are given a target node target and an integer k. Return a list of all the values that are at a distance k from the target node in the tree.

Approach

The idea to solve this problem is to perform a depth-first search (DFS) traversal of the binary tree to find nodes at a distance k from the target node. We'll use a few helper functions to achieve this.

1. UpdateParent Function: We'll start by creating a mapping of each node to its parent node. This will help us navigate the tree later.

2. Traversal Function: In this function, we'll perform the DFS traversal. When we reach a node at a distance k from the target node, we add its value to the result list. We also maintain a HashSet to keep track of visited nodes to avoid revisiting them. The recursion continues until we either find nodes at distance k or reach the end of the tree.

3. Finally, in the distanceK function, we initialize the necessary data structures and call the UpdateParent function to create the mapping of nodes to their parents. Then, we call the Traversal function to find nodes at distance k from the target node and add their values to the result list.

Complexity Analysis

• Time Complexity: The DFS traversal of the binary tree takes O(n) time, where n is the number of nodes in the tree.
• Space Complexity: The space complexity is O(n) due to the space used by the HashMap, HashSet, and the recursion stack.

Code :

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {

    public void UpdateParent(TreeNode Parent, TreeNode Current, HashMap<TreeNode, TreeNode> map) {
        if (Current == null) {
            return;
        }
        map.put(Current, Parent);

        UpdateParent(Current, Current.left, map);
        UpdateParent(Current, Current.right, map);
    }

    public void Traversal(int distance, TreeNode node, HashMap<TreeNode, TreeNode> map, List<Integer> list,HashSet<TreeNode> set) {
        if (node == null) {
            return;
        }
        if (set.contains(node)) {
            return;
        }
        set.add(node);

        if (distance == 0) {
            list.add(node.val);
            return;
        }
        Traversal(distance - 1, node.left, map, list, set);
        Traversal(distance - 1, node.right, map, list, set);
        Traversal(distance - 1, map.get(node), map, list, set);

    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        HashMap<TreeNode, TreeNode> map = new HashMap<>();
        HashSet<TreeNode> set = new HashSet<>();
        List<Integer> list = new ArrayList<>();
        UpdateParent(null, root, map);
        Traversal(k, target, map, list, set);
        return list;
    }
}