【剑指offer】JZ54 二叉搜索树的第k个节点

1 问题

数据范围：
0≤n≤1000，0≤k≤1000，树上每个结点的值满足0≤val≤1000
进阶：空间复杂度 O(n)，时间复杂度 O(n)

如输入{5,3,7,2,4,6,8},3时，二叉树{5,3,7,2,4,6,8}如下图所示：

该二叉树所有节点按结点值升序排列后可得[2,3,4,5,6,7,8]，所以第3个结点的结点值为4，故返回对应结点值为4的结点即可。

示例1
输入：{5,3,7,2,4,6,8},3
返回值：4

示例2
输入：{},1
返回值：-1

2 答案

自己写的，回溯算法

class Solution:
    def KthNode(self , proot: TreeNode, k: int) -> int:
        res = []
        self.dfs(proot, res)
        return res[k-1] if k-1 >= 0 and k-1 <= len(res) - 1 else -1
        
    def dfs(self, proot, res):
        if not proot:
            return 
        # if not proot.left and not proot.right:
        #     return 
        # 不能这样写，因为最后一个叶节点不会遍历
        self.dfs(proot.left, res)
        res.append(proot.val)
        self.dfs(proot.right, res)

官方解

1. 递归中序遍历

根据二叉搜索树的性质，左子树的元素都小于根节点，右子树的元素都大于根节点。因此它的中序遍历（左中右）序列正好是由小到大的次序，因此我们可以尝试递归中序遍历，也就是从最小的一个节点开始，找到k个就是我们要找的目标。

• step 1：设置全局变量count记录遍历了多少个节点，res记录第k个节点。
• step 2：另写一函数进行递归中序遍历，当节点为空或者超过k时，结束递归，返回。
• step 3：优先访问左子树，再访问根节点，访问时统计数字，等于k则找到。
• step 4：最后访问右子树。

class Solution:
    def __init__(self):
        self.res = None
        self.count = 0
        
    def midOrder(self, root, k):
        if not root or self.count > k:
            return 
        self.midOrder(root.left, k)
        self.count+=1
        if self.count == k:
            self.res = root # 不在这里return，在上面return
        self.midOrder(root.right, k)

    def KthNode(self , proot: TreeNode, k: int) -> int:
        self.midOrder(proot, k)
        if self.res:
            return self.res.val
        else:
            return -1

2. 非递归中序遍历，栈

递归实际上就是一种先进后出的栈结构，因此能用递归进行的中序遍历，非递归（栈）也可以实现，还是需要记录遍历到第k位即截止。

• step 1：用栈记录当前节点，不断往左深入，直到左边子树为空。
• step 2：再弹出栈顶（即为当前子树的父节点），访问该节点，同时计数。
• step 3：然后再访问其右子树，其中每棵子树都遵循左中右的次序。
• step 4：直到第k个节点返回，如果遍历结束也没找到，则返回-1.

class Solution:
    def KthNode(self , proot: TreeNode, k: int) -> int:
        if not proot:
            return -1
        count = 0
        p = None
        s = []
        while len(s) != 0 or proot is not None:
            while proot is not None:
                s.append(proot)
                proot = proot.left
            p = s[-1]
            s.pop()
            count += 1
            if count == k:
                return p.val
            proot = p.right
        return -1

https://www.nowcoder.com/share/jump/9318638301699766385619

也可以用颜色标记迭代法

class Solution:
    def KthNode(self , proot: TreeNode, k: int) -> int:
        WhITE, GRAY = 0, 1
        count = 0
        stack = [(WhITE, proot)]
        while stack:
            color, node = stack.pop()
            if node is None: continue
            if color == WhITE:
                stack.append((WhITE, node.right))
                stack.append((GRAY, node))
                stack.append((WhITE, node.left))
            else:
                # 这里才是真正有效遍历节点的地方
                count += 1
                if count == k:
                    return node.val
        return -1

https://blog.csdn.net/CSDNLHCC/article/details/134019651?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522169976648016800213014996%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fblog.%2522%257D&request_id=169976648016800213014996&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2_blogfirst_rank_ecpm_v1~rank_v31_ecpm-1-134019651-null-null.nonecase&utm_term=%E4%BA%8C%E5%8F%89%E6%A0%91&spm=1018.2226.3001.4450