距离是K的二叉树节点

# -*- 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