algoadvance

Leetcode 863. All Nodes Distance K in Binary Tree

Problem Statement

You are given the root of a binary tree, 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.

The distance between two nodes is defined by the number of edges on the path from one to the other.

Clarifying Questions

  1. Can the target node be the root node?
    • Yes, the target node can be any node in the tree, including the root node.
  2. Can the tree have duplicate values?
    • Yes, but each node is unique.
  3. What should be returned if there are no nodes at distance k?
    • Return an empty list if no nodes are at distance k.
  4. What is the maximum depth of the binary tree?
    • The constraints generally follow typical binary tree problems, so let’s assume a reasonable depth.

Strategy

  1. Build a Parent Reference Map:
    • We need to keep track of parent pointers for each node so that we can traverse upwards in the tree.
  2. Perform BFS from the Target Node:
    • Use a Breadth-First Search (BFS) strategy starting from the target node to find all nodes at distance k. Include traversal up to parent nodes using the previously built parent reference map.

Code

import java.util.*;

public class Solution {
    // TreeNode definition
    static class TreeNode {
        int val;
        TreeNode left, right;
        TreeNode(int x) { val = x; }
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        Map<TreeNode, TreeNode> parentMap = new HashMap<>();
        buildParentMap(root, null, parentMap);

        Queue<TreeNode> queue = new LinkedList<>();
        Set<TreeNode> visited = new HashSet<>();
        queue.add(target);
        visited.add(target);

        int currentLevel = 0;
        while (!queue.isEmpty()) {
            if (currentLevel == k) {
                List<Integer> result = new ArrayList<>();
                for (TreeNode node : queue) {
                    result.add(node.val);
                }
                return result;
            }

            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode currentNode = queue.poll();
                for (TreeNode neighbor : getNeighbors(currentNode, parentMap)) {
                    if (!visited.contains(neighbor)) {
                        visited.add(neighbor);
                        queue.add(neighbor);
                    }
                }
            }
            currentLevel++;
        }

        return new ArrayList<>();
    }

    private void buildParentMap(TreeNode node, TreeNode parent, Map<TreeNode, TreeNode> parentMap) {
        if (node != null) {
            parentMap.put(node, parent);
            buildParentMap(node.left, node, parentMap);
            buildParentMap(node.right, node, parentMap);
        }
    }

    private List<TreeNode> getNeighbors(TreeNode node, Map<TreeNode, TreeNode> parentMap) {
        List<TreeNode> neighbors = new ArrayList<>();
        if (node.left != null) neighbors.add(node.left);
        if (node.right != null) neighbors.add(node.right);
        if (parentMap.get(node) != null) neighbors.add(parentMap.get(node));
        return neighbors;
    }

    public static void main(String[] args) {
        // Example usage:
        Solution solution = new Solution();
        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);

        TreeNode target = root.left; // target is the node with value 5
        int k = 2;

        System.out.println(solution.distanceK(root, target, k)); // Output: [7, 4, 1]
    }
}

Time Complexity

Cut your prep time in half and DOMINATE your interview with AlgoAdvance AI