CLOSE
🛠️ Settings

Problem Statement

Given the root of a binary tree, a target node, and an integer K, return the list of the values of all nodes that are exactly K distance away from the target node. The distance between two nodes is defined as the number of edges in the shortest path connecting them.

Input:

  • root: The root node of the binary tree.
  • target: The target node in the binary tree.
  • K: The distance from the target node.

Output:

  • A list of integers representing the values of nodes that are at distance KKK from the target node. The order of values in the output list does not matter.

Examples

Example 1:

Input: Target node = 5, K = 2
        3
       / \
      5   1
     / \  / \
    6   2 0  8
       /   \
      7     4
Output: [7, 4, 1]

Explanation:
Nodes at distance 2 from node 5 are:
  Node 7 (child of node 2)
  Node 4 (child of node 2)
  Node 1 (child of node 3)
Example 2:

Input: Target node = 1, K = 1
        1
       / \
      2   3
Output: [2, 3]

Explanation:
Nodes at distance 1 from node 1 are 2 and 3.
Example 3:

Input: Target node = 1, K = 0
        1
Output = [1]

Explanation:
Node at distance 0 from itself is 1.

Different Approaches

1️⃣ 

Intuition:

To tackle the problem of finding all nodes at a specific distance 'K' from a given node in a binary tree, we need to understand the spatial distribution of nodes around the target node. Nodes that are exactly 'K' steps away from the target node can be thought of as forming a concentric pattern around it, with each level of distance increasing by one step. To efficiently traverse the tree and find these nodes, access not only the direct child nodes but also the parent nodes. While child nodes can be accessed via pointers, the parent node access requires constructing a mapping from each node to its parent. This involves three main phases: constructing the parent-child relationships using a breadth-first search (BFS), identifying and storing the target node, and then performing a depth-first search (DFS) to explore nodes at the required distance 'K' from the target node.

Approach:

  1. Begin by initializing a queue and a hashmap to keep track of parent nodes. Start by inserting the root node into the queue. Process each node by recording its parent and enqueuing its children. This process builds a comprehensive map of parent-child relationships throughout the tree.
  2. If a reference to the target node is provided directly, use that reference. If only the target node’s value is given, traverse the tree to find and store the node with this value, ensuring you have a reference to it for the next steps.
  3. Initialize a queue and a hashmap for tracking visited nodes. Insert the target node into the queue with an initial distance of 0. While processing nodes in the queue, examine the current node’s adjacent nodes (left child, right child, and parent). For each adjacent node, if it hasn't been visited yet, mark it as visited, enqueue it with an updated distance, and continue the traversal. When the distance 'dis' reaches 'K', add the node to the result list.
  4. After completing the DFS traversal, return the list of nodes that are exactly 'K' steps away from the target node.

Code:

#include <iostream>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <queue>
using namespace std;

// Definition of a TreeNode
struct TreeNode {
    int data;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int val) : data(val), left(nullptr), right(nullptr) {}
};

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        vector<int> result;

        // Step 1: Create a parent map
        unordered_map<TreeNode*, TreeNode*> parentMap;
        populateParentMap(root, nullptr, parentMap);

        // Step 2: Perform BFS from the target node
        unordered_set<TreeNode*> visited;
        queue<pair<TreeNode*, int>> q;
        q.push({target, 0});
        visited.insert(target);

        while (!q.empty()) {
            auto [node, distance] = q.front();
            q.pop();

            // If we reach distance K, collect all nodes at this level
            if (distance == k) {
                result.push_back(node->data);
            }

            // Explore neighbors (left, right, and parent)
            if (node->left && visited.find(node->left) == visited.end()) {
                q.push({node->left, distance + 1});
                visited.insert(node->left);
            }
            if (node->right && visited.find(node->right) == visited.end()) {
                q.push({node->right, distance + 1});
                visited.insert(node->right);
            }
            if (parentMap[node] && visited.find(parentMap[node]) == visited.end()) {
                q.push({parentMap[node], distance + 1});
                visited.insert(parentMap[node]);
            }
        }

        return result;
    }

private:
    // Helper function to populate parent map
    void populateParentMap(TreeNode* node, TreeNode* parent, unordered_map<TreeNode*, TreeNode*>& parentMap) {
        if (!node) return;
        if (parent) parentMap[node] = parent;
        populateParentMap(node->left, node, parentMap);
        populateParentMap(node->right, node, parentMap);
    }
};

// Helper function to test the solution
void testSolution() {
    TreeNode* root = new TreeNode(3);
    root->left = new TreeNode(5);
    root->right = new TreeNode(1);
    root->left->left = new TreeNode(6);
    root->left->right = new TreeNode(2);
    root->right->left = new TreeNode(0);
    root->right->right = new TreeNode(8);
    root->left->right->left = new TreeNode(7);
    root->left->right->right = new TreeNode(4);

    Solution sol;
    vector<int> result = sol.distanceK(root, root->left, 2); // Target node = 5, K = 2
    for (int val : result) {
        cout << val << " ";
    }
    cout << endl;
}

int main() {
    testSolution();
    return 0;
}

Complexity Analysis:

  • Time Complexity: O(N)
    • Parent Mapping: O(N)
    • BFS: O(N), in the worst case, as each node is visited once.
    • Overall: O(N)
  • Space Complexity: O(N)
    • Parent map: O(N)
    • Queue for BFS: O(N) in the worst case.
    • Visited set: O(N)