【LeetCode 230】二叉搜索树中第K小的元素

1.题目描述

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1

   3
  / \
 1   4
  \
   2

输出: 1
示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3

       5
      / \
     3   6
    / \
   2   4
  /
 1

输出: 3
进阶：
如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化 kthSmallest 函数？

2.第一种方法

看见这题脑子里蹦出的第一个想法就是中序遍历，因为二叉搜索树的中序遍历是升序序列，得到该序列后直接定位到第 k - 1 个元素就是第 k 小的元素。于是就有了下面的代码：

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        result = []
        self.inorder(root, result)
        return result[k-1]

    def inorder(self, root, result):
        if not root:
            return
        if root.left:
            self.inorder(root.left, result)
        result.append(root.val)
        if root.right:
            self.inorder(root.right, result)

也可以在递归里加上返回值，上面的做一下改动，像下面这样：

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        result = []
        res = self.inorder(root, result)   # 不一样的地方
        return res[k-1]

    def inorder(self, root, result):
        if not root:
            return
        if root.left:
            self.inorder(root.left, result)
        result.append(root.val)
        if root.right:
            self.inorder(root.right, result)
        return result

只是这种方法着实太慢，翻看别人的题解才知道问题出在哪，我这种相当于把整个树遍历完了再去定位节点，若指定的第K小元素刚好是中序遍历的第一个元素，那应该找到了就返回，后面的就不用再遍历了，上面的代码就相当于做了无用功。所以改进如下：

3.第二种方法

思路： 采取迭代式的中序遍历，并初始化一个计数器 count 计算当前遍历到哪个节点，当 count = k 时，表明已经找到，直接返回当前节点值即可。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        if not root:
            return 
        count = 0
        stack = []
        node = root
        while stack or node:
            while node:
                stack.append(node)
                node = node.left
            node = stack.pop()
            count += 1
            if count == k:
                return node.val
            node = node.right

4.第三种方法

思路: 采取迭代式中序遍历，跟第二种思路差不多，只不过换了一个列表存储遍历过的元素。当元素个数等于k时停止遍历，最后一个就是要找的元素。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        if root is None:
            return root     
        stack = []
        container = []
        curr = root
        while curr or stack:
            while curr:
                stack.append(curr)
                curr = curr.left
            node = stack.pop()
            container.append(node.val)
            if len(container) == k:
                break
            curr = node.right
        return container.pop()