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

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:returninorder(node.left)tree.append(node.val)inorder(node.right)return treeinorder(root)return tree[k-1]

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

优化1

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

class Solution(object):def kthSmallest(self, root, k):self.count = 0self.result = Nonedef inorder(node):if not node or self.result is not None:returninorder(node.left)self.count += 1if self.count == k:self.result = node.valreturninorder(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.leftroot = stack.pop()k -= 1if k == 0:return root.valroot = root.right

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