LC 863: All Nodes Distance K in Binary Tree

Difficulty: Medium
Category: Tree, DFS, BFS
Companies: Amazon, Facebook, Google, Microsoft, Apple

Problem Statement

Given the root of a binary tree, the value of a target node, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

Examples

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output: [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (value 5) have values 7, 4, and 1.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

Constraints

  • The number of nodes in the tree is in the range [1, 500].
  • 0 <= Node.val <= 500
  • All the values Node.val are unique.
  • target is the value of one of the nodes in the tree.
  • k >= 0

Solution Approaches

Approach 1: Convert Tree to Graph with DFS

Key Insight: Convert the binary tree into an undirected graph, then perform BFS/DFS from the target node to find all nodes at distance k.

Algorithm:

  1. Build adjacency list by traversing the tree
  2. Store parent-child relationships bidirectionally
  3. Perform DFS starting from target node
  4. Track visited nodes to avoid cycles
  5. Collect nodes at exact distance k

Time Complexity: O(n)
Space Complexity: O(n)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        buildGraph(root, nullptr);
        visited.insert(target->val);
        dfs(target->val, 0, k);
        return rtn;
    }
private:
    unordered_map<int, vector<int>> graph;
    vector<int> rtn;
    unordered_set<int> visited;

    void buildGraph(TreeNode* curr, TreeNode* parent) {
        if(curr && parent) {
            graph[curr->val].push_back(parent->val);
            graph[parent->val].push_back(curr->val);
        }
        if(curr->left) buildGraph(curr->left, curr);
        if(curr->right) buildGraph(curr->right, curr);
    }

    void dfs(int curr, int dist, int K) {
        if(dist == K) {
            rtn.push_back(curr);
            return;
        }
        for(int neighbor: graph[curr]) {
            if(!visited.contains(neighbor)) {
                visited.insert(neighbor);
                dfs(neighbor, dist + 1, K);
            }
        }
    }
};

Approach 2: Parent Pointer Map with DFS

Key Insight: Create a parent pointer map to traverse up the tree, then use DFS to explore all neighbors (left, right, parent).

Algorithm:

  1. Traverse tree to build parent map
  2. Start DFS from target node
  3. For each node, explore left, right, and parent
  4. Track visited nodes
  5. Collect nodes at distance k

Time Complexity: O(n)
Space Complexity: O(n)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        addParent(root, nullptr);
        vector<int> rtn;
        unordered_set<TreeNode*> visited;
        dfs(target, k, rtn, visited);
        return rtn;
    }

private:
    unordered_map<TreeNode*, TreeNode*> parent;
    
    void addParent(TreeNode* curr, TreeNode* parent) {
        if(curr) {
            this->parent[curr] = parent;
            addParent(curr->left, curr);
            addParent(curr->right, curr);
        }
    }

    void dfs(TreeNode* curr, int dist, vector<int>& rtn, unordered_set<TreeNode*>& visited) {
        if(!curr || visited.contains(curr)) return;
        visited.insert(curr);
        if(dist == 0) {
            rtn.push_back(curr->val);
            return;
        }
        dfs(parent[curr], dist - 1, rtn, visited);
        dfs(curr->left, dist - 1, rtn, visited);
        dfs(curr->right, dist - 1, rtn, visited);
    }
};

Approach 3: BFS with Graph Representation

Algorithm:

  1. Build adjacency list representation
  2. Use BFS with queue to find nodes at distance k
  3. Track level/distance during traversal
  4. Collect nodes exactly at distance k

Time Complexity: O(n)
Space Complexity: O(n)

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        buildGraph(root, nullptr);
        vector<int> rtn;
        queue<pair<int, int>> q;  // {node, distance}
        unordered_set<int> visited;
        
        q.push({target->val, 0});
        visited.insert(target->val);
        
        while(!q.empty()) {
            auto [curr, dist] = q.front();
            q.pop();
            
            if(dist == k) {
                rtn.push_back(curr);
            } else if(dist < k) {
                for(int neighbor: graph[curr]) {
                    if(!visited.contains(neighbor)) {
                        visited.insert(neighbor);
                        q.push({neighbor, dist + 1});
                    }
                }
            }
        }
        return rtn;
    }

private:
    unordered_map<int, vector<int>> graph;
    vector<int> rtn;
    
    void buildGraph(TreeNode* curr, TreeNode* parent) {
        if(curr && parent) {
            graph[curr->val].push_back(parent->val);
            graph[parent->val].push_back(curr->val);
        }
        if(curr->left) buildGraph(curr->left, curr);
        if(curr->right) buildGraph(curr->right, curr);
    }
};

Algorithm Analysis

Approach Comparison

Approach Time Space Pros Cons
Graph + DFS O(n) O(n) Simple, intuitive Extra memory for graph
Parent Map + DFS O(n) O(n) No extra graph structure More complex traversal
Graph + BFS O(n) O(n) Level-order traversal Slightly more code

Key Insights

  1. Tree as Graph: Binary trees are directed acyclic graphs; adding parent pointers makes them undirected
  2. Multi-directional Search: Can traverse up (parent) and down (children) in tree
  3. Cycle Prevention: Must track visited nodes to avoid infinite loops
  4. Distance Tracking: Depth-first or breadth-first approaches both work

Implementation Details

Building Parent-Child Relationships

// Store bidirectional edges
if(curr && parent) {
    graph[curr->val].push_back(parent->val);
    graph[parent->val].push_back(curr->val);
}

DFS with Distance Control

void dfs(int curr, int dist, int K) {
    if(dist == K) {
        rtn.push_back(curr);
        return;
    }
    // Recursively explore all neighbors
    for(int neighbor: graph[curr]) {
        if(!visited.contains(neighbor)) {
            visited.insert(neighbor);
            dfs(neighbor, dist + 1, K);
        }
    }
}

// Explore: parent, left child, right child
dfs(parent[curr], dist - 1, rtn, visited);
dfs(curr->left, dist - 1, rtn, visited);
dfs(curr->right, dist - 1, rtn, visited);

Edge Cases

  1. Target is Root: Still works, no parent path
  2. k = 0: Returns only target node
  3. Single Node Tree: Returns empty array if k > 0
  4. k Beyond Tree Depth: Returns empty array
  5. Target at Leaf: Must go up to parent then down

Follow-up Questions

  • What if the tree had more than 2 children per node?
  • How would you modify for a directed graph?
  • What if nodes could have duplicate values?
  • How would you find nodes within distance k (not exactly k)?

Optimization Techniques

  1. Graph Conversion: O(n) one-time preprocessing
  2. Visited Tracking: O(1) lookup prevents redundant work
  3. Early Termination: Stop at distance k exactly
  4. Bidirectional Edges: Store parent-child relationships for undirected graph behavior

Code Quality Notes

  1. Clarity: Graph approach is most intuitive for new learners
  2. Modularity: Separate graph building from search logic
  3. Memory: Both approaches use O(n) space for n nodes
  4. Performance: All solutions are optimal O(n) time

This problem beautifully demonstrates how binary trees can be treated as graphs when we need multi-directional traversal capabilities.