230. Kth Smallest Element in a BST -- python

最新推荐文章于 2021-04-09 14:18:22 发布

原创最新推荐文章于 2021-04-09 14:18:22 发布 · 502 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#leetcode

leetcode 专栏收录该内容

13 篇文章

订阅专栏

本文介绍了一种寻找二叉搜索树中第K小元素的方法，利用二叉搜索树的特性，通过中序遍历可以得到节点值升序排列的结果。文中提供了两种实现方式：迭代法和递归法。

题目：

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Credits:
Special thanks to @ts for adding this problem and creating all test cases.

解题：

因为左节点小于根节点小于右节点，二叉搜索树的一个特性就是中序遍历的结果就是树内节点从小到大顺序输出的结果。这里采用迭代形式，我们就可以在找到第k小节点时马上退出。

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        stack = []
        while root or stack:  #当root不等于none或stack不等于none
            while root:       # 搜索树的最左边的节点
                stack.append(root)
                root = root.left
            root = stack.pop()   
            k -= 1
            if k == 0 :
                return root.val
            root = root.right

递归形式：

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        count = []
        self.helper(root, count)
        return count[k-1]
        
    def helper(self, node, count):
        if not node:
            return
        self.helper(node.left, count)
        count.append(node.val)
        self.helper(node.right, count)