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 + " ");
        }
    }

}