Problem

Given a binary tree, return or print all the nodes which are at distance X from the root.

Examples

Example 1

1
2
3
4
5
6

Input: root = [1, 2, 3, 4, 5, null, null], X = 2
Output: [4, 5]
Explanation:
The nodes at a distance of 2 from the root node 1 are 4 and 5.
- Distance 2 from root 1 -> Node 4
- Distance 2 from root 1 -> Node 5

Example 2

1
2
3
4
5
6

Input: root = [1, 2, 3, null, 4, 5, null], X = 1
Output: [2, 3]
Explanation:
The nodes at a distance of 1 from the root node 1 are 2 and 3.
- Distance 1 from root 1 -> Node 2
- Distance 1 from root 1 -> Node 3

Similar Problem

Nodes at Distance k from Leaf Nodes in Binary Tree All Nodes Distance K in Binary Tree

Solution

To find all the nodes at distance X from the root, we can use Depth First Search (DFS) or Breadth First Search (BFS). The idea is to traverse the tree and keep track of the current depth. When the depth equals X, we collect the node’s value.

Method 1 - DFS

  • Traverse the tree using DFS. (pre-order traversal)
  • Keep track of the current depth of each node.
  • When the depth equals X, collect the node’s value.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

public class Solution {
    public List<Integer> nodesAtDistanceKFromRoot(TreeNode root, int k) {
        List<Integer> ans = new ArrayList<>();
        dfs(root, k, 0, ans);
        return ans;
    }

    private void dfs(TreeNode node, int k, int depth, List<Integer> ans) {
        if (node == null) {
            return;
        }

        if (depth == k) {
            ans.add(node.val);
            return;
        }

        dfs(node.left, k, depth + 1, ans);
        dfs(node.right, k, depth + 1, ans);
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution:
    def nodesAtDistanceKFromRoot(self, root: Optional[TreeNode], k: int) -> List[int]:
        def dfs(node: Optional[TreeNode], k: int, depth: int, ans: List[int]) -> None:
            if not node:
                return
            
            if depth == k:
                ans.append(node.val)
                return
            
            dfs(node.left, k, depth + 1, ans)
            dfs(node.right, k, depth + 1, ans)
        
        ans = []
        dfs(root, k, 0, ans)
        return ans

Complexity

  • ⏰ Time complexity: O(n), where n is the number of nodes in the tree, since each node is visited once.
  • 🧺 Space complexity: O(h), where h is the maximum depth of the tree for the recursion stack in DFS.