BST中第K小的元素

https://www.lintcode.com/problem/902/?utm_source=sc-cheatsheet-cyc

描述

给一棵二叉搜索树,写一个 KthSmallest 函数来找到其中第 K 小的元素。

你可以假设 k 总是有效的, 1 ≤ k ≤ 树的总节点数

样例

样例 1:

输入:{1,#,2},2输出:2解释:	1	 \	  2第二小的元素是2。

样例 2:

输入:{2,1,3},1输出:1解释:  2 / \1   3第一小的元素是1。

挑战

如果这棵 BST 经常会被修改(插入/删除操作)并且你需要很快速的找到第 K 小的元素呢?你会如何优化这个 KthSmallest 函数?

解题思路

中序简历,做K次就行。 注意while循环考虑stack是左树,cur是右树,两边子树的情况都要考虑。

```python
from lintcode import (
    TreeNode,
)

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None
"""

class Solution:
    """
    @param root: the given BST
    @param k: the given k
    @return: the kth smallest element in BST
    """
    def kth_smallest(self, root: TreeNode, k: int) -> int:
        # write your code here
        if not root:
            return None
        s = []
        cur = root
        res = None
        while k > 0 and (s or cur): #s or cur, 是因为stack存左树,cur存右树。所有左右树都没有的情况都要考虑
            while cur:
                s.append(cur)
                cur = cur.left
            cur = s.pop()
            k -= 1
            res = cur.val
            cur = cur.right
        return res


```

Last updated