距离是K的二叉树节点

import java.util.*;

// 以节点target 为根,       下面第k层节点
// 以节点target 父节点 为根, 下面第k-1层节点
// 父父节点, k-2层 ...  循环
//
// Map存父节点
public class Solution {
    ArrayList<Integer> res = new ArrayList<>();  // 结果
    Set<TreeNode> parent = new HashSet<>(); // target及父节点

    public ArrayList<Integer> distanceKnodes (TreeNode root, int target, int k) {
        // BFS: 找到target节点及其父节点
        Map<TreeNode, TreeNode> map = new HashMap<>(); // 存父节点
        map.put(root, null);
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        TreeNode targetNode = null; // target节点
        while (!queue.isEmpty()) {
            TreeNode t = queue.poll();
            if (t.val == target) {
                targetNode = t;
                break;
            }
            if (t.left  != null) {
                map.put(t.left, t);
                queue.offer(t.left);
            }
            if (t.right != null) {
                map.put(t.right, t);
                queue.offer(t.right);
            }
        }
        // target及父节点
        TreeNode p = targetNode;
        while (p != null) {
            parent.add(p);
            p = map.get(p);
        }
        // 循环：target节点为根 下面k层节点, 父 k-1, 父父 k-2...
        if (k == 0) res.add(target); // k==0 特判target
        p = targetNode;
        while (p != null) {
            bfs(p, k, target); // 调用
            p = map.get(p);
            k--;
        }
        
        return res;
    }

    // BFS: 以root为根, 下面第k层节点 (排除target及其父节点那侧子树🍓)
    void bfs(TreeNode root, int k, int target) {
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        int level = 0;
        while (!queue.isEmpty() && level <= k + 1) {
            level++;
            int size = queue.size();
            while (size-- > 0) {
                TreeNode t = queue.poll();
                if (level == k + 1) { // 下面第k层节点
                    res.add(t.val);
                }
                // 排除target及其父节点 那个子树分支（因为会重复计算,上去时已算过一部分路径了）
                if (t.left  != null && !parent.contains(t.left))  queue.offer(t.left);
                if (t.right != null && !parent.contains(t.right)) queue.offer(t.right);
            }
        }
    }
}

import java.util.*;

/*
 * public class TreeNode {
 *   int val = 0;
 *   TreeNode left = null;
 *   TreeNode right = null;
 *   public TreeNode(int val) {
 *     this.val = val;
 *   }
 * }
 */

/**
 * NC370 距离是K的二叉树节点
 * @author d3y1
 */
public class Solution {
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     *
     * @param root TreeNode类
     * @param target int整型
     * @param k int整型
     * @return int整型一维数组
     */
    public ArrayList<Integer> distanceKnodes (TreeNode root, int target, int k) {
        // return solution1(root, target, k);
        return solution2(root, target, k);
    }

    private HashSet<Integer>[] adj;
    private boolean[] isVisited;
    private ArrayList<Integer> result = new ArrayList<>();
    private final int N = 1000;

    /**
     * 后序遍历 + 层序遍历
     * 
     * 图模型
     * 
     * @param root
     * @param target
     * @param k
     * @return
     */
    private ArrayList<Integer> solution1(TreeNode root, int target, int k){
        isVisited = new boolean[N];
        adj = new HashSet[N];
        for(int i = 0; i < N; i++){
            adj[i] = new HashSet<>();
        }

        postOrder(root);
        levelOrder(target, k);

        return result;
    }

    /**
     * 递归: 后序遍历
     * @param root
     */
    private void postOrder(TreeNode root){
        if(root == null){
            return;
        }

        postOrder(root.left);
        postOrder(root.right);

        if(root.left != null){
            adj[root.val].add(root.left.val);
            adj[root.left.val].add(root.val);
        }
        if(root.right != null){
            adj[root.val].add(root.right.val);
            adj[root.right.val].add(root.val);
        }
    }

    /**
     * 队列: 层序遍历
     * @param target
     * @param k
     */
    private void levelOrder(int target, int k){
        Queue<Integer> queue = new LinkedList<>();
        isVisited[target] = true;
        queue.offer(target);

        int curr;
        int size;
        int level = 0;
        boolean isFound = false;
        while(!queue.isEmpty()){
            level++;
            size = queue.size();
            while(size-- > 0){
                curr = queue.poll();
                for(int next: adj[curr]){
                    if(!isVisited[next]){
                        isVisited[next] = true;
                        if(level == k){
                            isFound = true;
                            result.add(next);
                        }else{
                            queue.offer(next);
                        }
                    }
                }
            }
            if(isFound){
                break;
            }
        }
    }

    //////////////////////////////////////////////////////////////////////////////////////////

    private ArrayList<Integer> list = new ArrayList<>();
    private HashMap<Integer,TreeNode> parentMap = new HashMap<>();

    /**
     * 层序遍历 + 前序遍历
     * @param root
     * @param target
     * @param k
     * @return
     */
    private ArrayList<Integer> solution2(TreeNode root, int target, int k){
        if(root == null){
            return list;
        }

        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        parentMap.put(root.val, null);


        TreeNode node,cNode,pNode;
        while(!queue.isEmpty()){
            node = queue.poll();
            if(node.val == target){
                if(k == 0){
                    list.add(node.val);
                    return list;
                }
                downKnodes(node, k);
                cNode = node;
                pNode = parentMap.get(cNode.val);
                while(pNode!=null && k>0){
                    upKnodes(cNode, pNode, --k);
                    cNode = pNode;
                    pNode = parentMap.get(cNode.val);
                }

                return list;
            }
            if(node.left != null){
                parentMap.put(node.left.val, node);
                queue.offer(node.left);
            }
            if(node.right != null){
                parentMap.put(node.right.val, node);
                queue.offer(node.right);
            }
        }

        return list;
    }

    /**
     * 向上遍历
     * @param cNode
     * @param pNode
     * @param k
     */
    private void upKnodes(TreeNode cNode, TreeNode pNode, int k){
        if(k == 0){
            list.add(pNode.val);
            return;
        }

        if(pNode.left!=null && pNode.left.val!=cNode.val){
            downKnodes(pNode.left, k-1);
        }

        if(pNode.right!=null && pNode.right.val!=cNode.val){
            downKnodes(pNode.right, k-1);
        }
    }

    /**
     * 向下遍历
     * @param root
     * @param k
     */
    private void downKnodes(TreeNode root, int k){
        preOrder(root, k);
    }

    /**
     * 前序遍历
     * @param root
     * @param k
     */
    private void preOrder(TreeNode root, int k){
        if(root == null){
            return;
        }

        if(k == 0){
            list.add(root.val);
            return;
        }

        preOrder(root.left, k-1);
        preOrder(root.right, k-1);
    }
}

class Solution {
private:
    unordered_map<int, TreeNode*> parents;
    unordered_set<TreeNode*> cnt;
    vector<int> result;
    // 找到所有节点的父节点
    void find_parents(TreeNode* root) {
        if (root==NULL) return;
        if (root->left) {
            parents[root->left->val]=root;
            find_parents(root->left);
        }
        if (root->right) {
            parents[root->right->val]=root;
            find_parents(root->right);
        }
    }
    // 层序遍历
    TreeNode* find_target(TreeNode* root, int target) {
        queue<TreeNode*> que;
        if (root==NULL) return root;
        que.push(root);
        while (!que.empty()) {
            int size=que.size();
            for (int i=0; i<size; i++) {
                TreeNode* cur=que.front();
                que.pop();
                if (cur->val==target) return cur;
                if (cur->left) que.push(cur->left);
                if (cur->right) que.push(cur->right);
            }
        }
        return NULL;
    }
    // 回溯三个方向求result
    void find_result(TreeNode* target, int& depth, int k) {
        if (target==0) return;
        if (depth==k) result.push_back(target->val);
        cnt.insert(target);
        if (cnt.count(target->left)==0) {
            depth++;
            find_result(target->left, depth, k);
            depth--;
        }
        if (cnt.count(target->right)==0) {
            depth++;
            find_result(target->right, depth, k);
            depth--;
        }
        if (cnt.count(parents[target->val])==0) {
            depth++;
            find_result(parents[target->val], depth, k);
            depth--;
        }
    }
public:
    vector<int> distanceKnodes(TreeNode* root, int target, int k) {
        find_parents(root);
        TreeNode* cur=find_target(root, target);
        int depth=0;
        find_result(cur, depth, k);
        return result;
    }
};

# -*- coding: utf-8 -*-


class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


class BT:
    def __init__(self):
        self.root = None

    def levelOrderToBT(self, nums):  # 层序遍历构造二叉树
        n = len(nums)

        if n > 0:
            self.root = TreeNode(nums[0])
            queue, idx = [self.root], 1

            while queue and idx < n:
                node = queue.pop(0)
                node.left = TreeNode(nums[idx]) if nums[idx] != None else None
                if node.left:
                    queue.append(node.left)
                node.right = TreeNode(nums[idx + 1]) \
                    if idx + 1 < n and nums[idx + 1] != None else None
                if node.right:
                    queue.append(node.right)
                idx += 2

        return self.root

    @staticmethod
    def levelOrder(root):  # 二叉树的层序遍历
        if not root:
            return []

        queue, res = [root], []
        while queue:
            node = queue.pop(0)
            if node:
                res.append(node.val)
                queue.append(node.left)
                queue.append(node.right)
            else:
                res.append(None)
        while res and res[-1] == None:
            res.pop()

        return res


#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
#
# @param root TreeNode类
# @param target int整型
# @param k int整型
# @return int整型一维数组
#
class Solution:
    """
    题目：
        https://www.nowcoder.com/practice/e280b9b5aabd42c9b36831e522485622?tpId=196&tqId=40501&rp=1&ru=/exam/oj&qru=/exam/oj&sourceUrl=%2Fexam%2Foj%3FjudgeStatus%3D3%26page%3D1%26pageSize%3D50%26search%3D%26tab%3D%25E7%25AE%2597%25E6%25B3%2595%25E7%25AF%2587%26topicId%3D196&difficulty=undefined&judgeStatus=3&tags=&title=
    参考：
        大神：姐姐的遮阳伞
    算法：
        1. 一次先序遍历二叉树，在寻找targetNode的同时，建立哈希表nodeMap，键：当前节点，值：该节点的父亲节点
        2. 对于距离为k的分析：
            a. 我们可以从targetNode出发，在其左右子树上寻找距离为k的节点
            b. 也可以通过哈希表，向上回溯到父亲节点进行搜索，只不过搜索的过程中，我们需要修改距离k：
                i) 回溯到targetNode的parent节点parentNode，若targetNode为parentNode的左孩子，那么我们从parentNode.right子树上寻找与根节点距离为k-2的节点；否则，若targetNode为parentNode的右孩子，那么我们从parentNode.left子树上寻找与根节点距离为k-2的节点
                ii) 再回溯到parentNode的父亲节点。继续搜索
                ...
                直到回溯到根节点或者k小于0为止
    复杂度：
        时间复杂度：O(N + kM) = O(k*N)，k为距离，M为每次搜索的平均节点数，N为二叉树的节点数
        空间复杂度：O(H), H为树的高度
    """

    def distanceKnodes(self, root, target, k):
        # write code here
        def helper(root):
            if root:
                if root.val == target:
                    self.targetNode = root
                if root.left:
                    nodeMap[root.left.val] = root
                if root.right:
                    nodeMap[root.right.val] = root
                helper(root.left)
                helper(root.right)

        def findAnsFromChild(node, k):  # 从node的左右子树中，寻找与node距离为k的节点
            if node:
                if k == 0:
                    if node.val != target:  # 搜索结果，剔除target
                        res.append(node.val)
                    return
                findAnsFromChild(node.left, k - 1)  # 递归搜索左子树
                findAnsFromChild(node.right, k - 1)  # 递归搜索右子树

        nodeMap, self.targetNode = {root.val: None}, None  # nodeMap记录节点node与其父亲节点的映射关系，targetNode为目标节点
        helper(root)

        res, startNode = [], self.targetNode
        findAnsFromChild(startNode, k)

        k -= 1  # 回溯到targetNode的父亲节点parentNode
        while k >= 0:
            parentNode = nodeMap[startNode.val]
            if not parentNode:  # 所有的父亲节点已经搜索完毕
                break

            if k == 0:  # 当回溯到根节点的距离恰好是k时，就无须探索了，因为再回溯，必然使得距离大于k
                res.append(parentNode.val)
                break

            if startNode == parentNode.left and k - 1 >= 0:  # startNode为parentNode的左孩子，就去parentNode的右子树搜索
                findAnsFromChild(parentNode.right, k - 1)
            if startNode == parentNode.right and k - 1 >= 0:  # startNode为parentNode的右孩子，就去parentNode的左子树搜索
                findAnsFromChild(parentNode.left, k - 1)

            startNode = parentNode  # 回溯到父亲节点
            k -= 1  # 距离-1

        return res


if __name__ == "__main__":
    sol = Solution()

    nums, target, k = [3, 5, 2, 4, 6, 0, 7, 1, 8], 5, 2

    # nums, target, k = [1, 2, 3, 4], 2, 1

    bt = BT()
    bt1 = bt.levelOrderToBT(nums)
    print BT.levelOrder(bt1)

    res = sol.distanceKnodes(bt1, target, k)
    print res

import java.util.*;

/*
 * public class TreeNode {
 *   int val = 0;
 *   TreeNode left = null;
 *   TreeNode right = null;
 *   public TreeNode(int val) {
 *     this.val = val;
 *   }
 * }
 */

public class Solution {
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param root TreeNode类 
     * @param target int整型 
     * @param k int整型 
     * @return int整型ArrayList
     */
    public ArrayList<Integer> distanceKnodes(TreeNode root, int target, int k) {
        // write code here
        if (root == null) {
            return new ArrayList<>();
        }
        // 最终返回的结果
        Set<Integer> finalRes = new HashSet<>();
        // 记录dfs遍历的所有路径
        ArrayList<ArrayList<Integer>> result = new ArrayList<>();
        // dfs遍历路径
        dfs(root, new ArrayList<>(), result);

        // 记录已经处理过的节点，理论上后续处理过的节点所在的路径无需再次处理
        ArrayList<Integer> processNode = new ArrayList<>();
        int father = target;
        while (k > 0) {
            judgeNode(finalRes, k--, result, processNode, father);
            // 加入到已处理集合
            processNode.add(father);
            // 找到新的父节点
            father = getFather(result, father);
        }
        return new ArrayList<>(finalRes);
    }

    private boolean alreadyChoose(ArrayList<Integer> processNodes, ArrayList<Integer> integers) {
        for (Integer processNode : processNodes) {
            if (integers.contains(processNode)) {
                return true;
            }
        }
        return false;
    }

    private void judgeNode(Set<Integer> finalRes, int k, ArrayList<ArrayList<Integer>> result,
                           ArrayList<Integer> processFather, int father) {
        for (ArrayList<Integer> integers : result) {
            if (alreadyChoose(processFather, integers)) {
                continue;
            }
            // 根节点，仅需处理未处理过的孩子
            if (father == -1) {
                if (k < integers.size()) {
                    finalRes.add(integers.get(k));
                }
            } else {
                // 非根节点，需处理包含father节点的路径
                if (integers.contains(father)) {
                    int index = integers.indexOf(father);
                    if (index + k < integers.size()) {
                        finalRes.add(integers.get(index + k));
                    }
                    if (index - k >= 0) {
                        finalRes.add(integers.get(index - k));
                    }
                }
            }

        }
    }

    /**
     * 获取当前处理节点的父节点
     * @param result
     * @param target
     * @return
     */
    private int getFather(ArrayList<ArrayList<Integer>> result, int target) {
        for (ArrayList<Integer> res : result) {
            if (res.contains(target)) {
                int index = res.indexOf(target);
                if (index - 1 >= 0) {
                    return res.get(index - 1);
                }
            }
        }
        // 已经是树根了，无父节点
        return -1;
    }

    private void dfs(TreeNode root, ArrayList<Integer> list, ArrayList<ArrayList<Integer>> result) {
        list.add(root.val);

        if (root.left == null && root.right == null) {
            result.add(new ArrayList<>(list));
            return;
        }

        if (root.left != null) {
            dfs(root.left, list, result);
            list.remove(list.size() - 1);
        }
        if (root.right != null) {
            dfs(root.right, list, result);
            list.remove(list.size() - 1);
        }
    }
}