230. 二叉搜索树中第 K 小的元素

• 给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 小的元素（从 1 开始计数）。

• 示例 1：

  • 输入：root = [3,1,4,null,2], k = 1
  • 输出：1

• 示例 1：

  • 输入：root = [5,3,6,2,4,null,null,1], k = 3
  • 输出：3

思路

由于二叉搜索树的中序遍历本身就是升序的，所以可以先将二叉树进行中序遍历，得到升序数组后直接输出第k-1个数即可

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: Optional[TreeNode]
        :type k: int
        :rtype: int
        """
        tree = []
        def inorder(node):  # 中序遍历
            if not node:
                return
            inorder(node.left)
            tree.append(node.val)
            inorder(node.right)
            return tree
        inorder(root)
        return tree[k-1]

• 时间复杂度: O(nlogn)
• 空间复杂度: O(n)

优化1

提前终止遍历（空间优化）
当找到第 k小的元素后，可以立即停止遍历：

class Solution(object):
    def kthSmallest(self, root, k):
        self.count = 0
        self.result = None
        def inorder(node):
            if not node or self.result is not None:
                return
            inorder(node.left)
            self.count += 1
            if self.count == k:
                self.result = node.val
                return
            inorder(node.right)
        inorder(root)
        return self.result

• 时间复杂度: O(k)，(最坏O(n))
• 空间复杂度: O(1)，(忽略递归栈空间，最坏O(n))

优化2

无递归栈溢出风险。
提前终止遍历，时间和空间最优。

class Solution(object):
    def kthSmallest(self, root, k):
        stack = []
        while root or stack:
            while root:
                stack.append(root)
                root = root.left
            root = stack.pop()
            k -= 1
            if k == 0:
                return root.val
            root = root.right

• 时间复杂度: O(k)
• 空间复杂度: O(h)，h是树高