package com.binaryTree;

import java.util.*;

// 题目描述:
// 给定三个参数:
// 二叉树的头结点head,树上某个节点target,正数K,
// 从target开始,可以向上走或者向下走,
// 返回与target的距离是K的所有节点
public class DistanceKNodes {

 public static class Node {
 public int value;
 public Node left;
 public Node right;

 public Node(int v) {
 value = v;
 }
 }

// 因为在经典二叉树中,当前节点不能向父节点方向走,就不能看到整棵树的全貌,就不能向上找到与当前节点 距离为K的节点
// 向上找到父节点的功能 可以用HashMap实现 parents

// 关键解题点:1、二叉树中 找到某个节点的父亲节点 使用HashMap实现,key为当前节点,value为父亲节点
// 2、通过queue.size() 来标记同一层/同一批/同一距离的 节点,


 public static List<Node> distanceKNodes(Node root, Node target, int K) {
 HashMap<Node,Node> father = new HashMap<>();
 father.put(root,null);
 generateFatherMap(root,father);

 // visited :节点元素是否进入过队列,
 Queue<Node> queue = new LinkedList<>();
 HashSet<Node> isVisited = new HashSet<>();

 queue.offer(target);
 isVisited.add(target);
 // curLevel : target 到其他节点的层数(路径)
 int curLevel = 0;
 // 距离为K的节点 收集在ans
 List<Node> res = new ArrayList<>();
 while (!queue.isEmpty()){
 int size = queue.size();
 while (size-- > 0){
 Node cur = queue.poll();
 if (curLevel == K){
 res.add(cur);
 }
 if (cur.left != null && !isVisited.contains(cur.left)){
 queue.offer(cur.left);
 isVisited.add(cur.left);
 }
 if (cur.right != null && !isVisited.contains(cur.right)){
 queue.offer(cur.right);
 isVisited.add(cur.right);
 }
 if (father.get(cur) != null && !isVisited.contains(father.get(cur))){
 queue.offer(father.get(cur));
 isVisited.add(father.get(cur));
 }
 }
 curLevel++;
 if (curLevel > K){
 break;
 }
 }
 return res;

 }

 public static void generateFatherMap(Node cur,HashMap<Node,Node> father){
 if (cur == null){
 return;
 }
 if (cur.left != null){
 father.put(cur.left,cur);
 generateFatherMap(cur.left,father);
 }

 if (cur.right != null){
 father.put(cur.right,cur);
 generateFatherMap(cur.right,father);
 }

 }

 public static void main(String[] args) {
 Node n0 = new Node(0);
 Node n1 = new Node(1);
 Node n2 = new Node(2);
 Node n3 = new Node(3);
 Node n4 = new Node(4);
 Node n5 = new Node(5);
 Node n6 = new Node(6);
 Node n7 = new Node(7);
 Node n8 = new Node(8);


 Node n11 = new Node(11);
 Node n12 = new Node(12);
 Node n13 = new Node(13);
 Node n14 = new Node(14);

 n3.left = n5;
 n3.right = n1;
 n5.left = n6;
 n5.right = n2;
 n1.left = n0;
 n1.right = n8;
 n2.left = n7;
 n2.right = n4;

 n4.left = n13;
 n4.right = n14;

 n8.left = n11;
 n8.right = n12;

 Node root = n3;
 Node target = n5;
 int K = 2;

 List<Node> ans = distanceKNodes(root, target, K);
 for (Node o1 : ans) {
 System.out.print(o1.value + " ");
 }
 System.out.println();

 K = 3;
 ans = distanceKNodes(root, target, K);
 for (Node o1 : ans) {
 System.out.print(o1.value + " ");
 }
 }

}