All Nodes Distance K in Binary Tree Problem

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 (with value 5) have values 7, 4, and 1.

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

public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
	Map <TreeNode, List<TreeNode>> g = new HashMap();
	List<Integer> ans = new ArrayList<Integer>();

	if (root == null || k <0) {
		return ans;
	}

	buildGraph(g, root, null);

	if (!g.containsKey(target)) {
		return ans;
	}

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

	while (!q.isEmpty()) {
		int size = q.size();

		if (k == 0) {
			for (int i = 0; i <size; i++) {
				ans.add(q.poll().val);
			}

			return ans;
		}

		for (int i = 0; i <size; i++) {
			TreeNode node = q.poll();

			for (TreeNode nei: g.get(node)) {
				if (visited.contains(nei)){
					 continue;
				}
				visited.add(nei);
				q.add(nei);
			}
		}

		k--;
	}

	return ans;
}


private void buildGraph(Map <TreeNode, List<TreeNode>> g, TreeNode node, TreeNode parent) {
	if (node == null) {
		return;
	}

	if (!g.containsKey(node)) {
		g.put(node, new ArrayList<TreeNode>());

		if (parent != null) {
			g.get(node).add(parent);
			g.get(parent).add(node);
		}

		buildGraph(g, node.left, node);
		buildGraph(g, node.right, node);
	}
}

public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
	List<Integer> ans = new ArrayList<>();
	Map<TreeNode,TreeNode> parentsMap = new HashMap<>();
	buildParentMap(root,root,parentsMap);//first root is node and second root is its parent considered
	Set<TreeNode> visited = new HashSet<>(); //keep track of visited nodes
	Queue<TreeNode> q = new LinkedList<>();//for BFS
	q.add(target);  //instead of root here we are adding target as we need to find from that node
	visited.add(target); //also target node need to be visited first
	int level = 0; //as we are using level order traversal k == level then we find all the nodes in queue
	while(q.size()>0){
		int size = q.size();
		if(level == K) {
			return buildListFromQueue(q);//elements remaining in queue will be nodes at dist 
		}
		for(int i = 0; i < size; i++){                  //k from target
			TreeNode t = q.remove();
			if(t.left != null && !visited.contains(t.left)){
				q.add(t.left);
				visited.add(t.left);
			}
			if(t.right != null && !visited.contains(t.right)){
				q.add(t.right);
				visited.add(t.right);
			}
			TreeNode parent = parentsMap.get(t);
			if(!visited.contains(parent)){ //if parent is not visited 
				q.add(parent);//add parent to q
				visited.add(parent);   //mark it as visited
			}
		}
		level++; //we check levelwise and each level covers 1 unit distance 
	}
		return ans;
}
private List<Integer> buildListFromQueue(Queue<TreeNode> q){
	List<Integer> lst = new ArrayList<>();
	while(q.size()>0){
		lst.add(q.remove().val);
	}
	return lst;
}
private void buildParentMap(TreeNode root,TreeNode parent,Map<TreeNode,TreeNode> parentsMap){
	if(root == null) {
		return;
	}
		
	parentsMap.put(root, parent);
	buildParentMap(root.left,root, parentsMap); 
	buildParentMap(root.right,root, parentsMap);
}