230. 二叉搜索树中第K小的元素
题意
给定一个二叉搜索树,编写一个函数 kthSmallest
来查找其中第 k 个最小的元素。 你可以假设 k 总是有效的,1 ≤ k ≤ 二叉搜索树元素个数。
解题思路
中序遍历,利用Python3中提供的生成器方法;
中序遍历,判断存储结点值的数组是否到到k,则表明访问的一个结点就是第k个最小的元素;
先获取跟结点处于的位置(第几个最小的元素),如果它比k小,则从右子结点中找,如果它比k大,则从左子节点中找;
实现
class Solution:
def kthSmallest(self, root: TreeNode, k: int) -> int:
"""
利用了Python3中新增的生成器方法;
"""
def gen(r):
if r is not None:
yield from gen(r.left)
yield r.val
yield from gen(r.right)
it = gen(root)
for _ in range(k):
ans = next(it)
return ans
def kthSmallest(self, root, k):
"""
中序遍历,判断当前访问的结点是否是第k个最小的个数;
:type root: TreeNode
:type k: int
:rtype: int
"""
stack = []
def helper(node):
if not node:
return
helper(node.left)
if len(stack) == k:
return;
stack.append(node.val)
helper(node.right)
if not root:
return None
helper(root)
return stack[-1]
def kthSmallest(self, root, k):
"""
:type root: TreeNode
:type k: int
:rtype: int
"""
def helper(node):
if not node:
return 0
# 获取以node为跟结点的结点总个数
return helper(node.left) + helper(node.right) + 1
if not root:
return None
while True:
# number表示当前结点是第n小的元素
number = helper(root.left) + 1
# 每次判断当前结点处在第几小的位置
# 如果比k大,说明值的范围比当前结点要小,则继续从当前结点的左子结点中找;
# 如果比k小,说明值的范围比当前结点要大,则继续从当前结点的右子结点中找;
if number < k:
root = root.right
k -= number
elif number > k:
root = root.left
else:
return root.val
return 0