[leetcode笔记] 230. 二叉搜索树中第K小的元素

题目描述

给定一个二叉搜索树，编写一个函数 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

解题思路

• 采取中序遍历的方式
  1. 通过递归的方法
  2. 通过循环的方法
• 还有网友写的另一种python独特的方法，就是通过 yeild 递归实现

代码

递归

# 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:
    	def zx(root):
            if not root:
                return []
            return zx(root.left)+[root.val]+zx(root.right)
        return zx(root)[k-1]

循环

# 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:
        i = 0
        stack = []
        curr = root
        while stack or curr:
            if curr:
                stack.append(curr)
                curr = curr.left
            else:
                curr = stack.pop()
                i += 1
                if i==k:
                    return curr.val
                curr = curr.right
        return 0

yeild递归实现

# 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:
    	def bt(node):
            if node is not None:
                yield from bt(node.left)
                yield node.val
                yield from bt(node.right)   
        i = bt(root)
        for _ in range(k):
            res = next(i)
        return res