k distance

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

 import java.util.*;

 class Solution {
     Map<TreeNode, TreeNode> parentMap = new HashMap<>();

     private void markParents(TreeNode root) {
         if (root == null) return;

         if (root.left != null) {
             parentMap.put(root.left, root);
             markParents(root.left);
         }
         if (root.right != null) {
             parentMap.put(root.right, root);
             markParents(root.right);
         }
     }
     private void bfs(TreeNode target, int k, List<Integer> result) {
         Queue<TreeNode> queue = new LinkedList<>();
         Set<TreeNode> visited = new HashSet<>();
         queue.offer(target);
         visited.add(target);
         while (!queue.isEmpty()) {
             int size = queue.size();
             if (k == 0) break;
             for (int i = 0; i < size; i++) {
                 TreeNode current = queue.poll();
                 if (current.left != null && !visited.contains(current.left)) {
                     queue.offer(current.left);
                     visited.add(current.left);
                 }
                 if (current.right != null && !visited.contains(current.right)) {
                     queue.offer(current.right);
                     visited.add(current.right);
                 }
                 if (parentMap.containsKey(current) && !visited.contains(parentMap.get(current))) {
                     queue.offer(parentMap.get(current));
                     visited.add(parentMap.get(current));
                 }
             }
             k--;
         }
         while (!queue.isEmpty()) {
             result.add(queue.poll().val);
         }
     }
     public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
         List<Integer> result = new ArrayList<>();
         // Step 1: Build parent mapping
         markParents(root);
         // Step 2: BFS from target node to find k-distance nodes
         bfs(target, k, result);
         return result;
     }
 }