距离是K的二叉树节点

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