leetcode-230_KthSmallestElementinaBST

上代码

#!/usr/bin/env python 
# encoding: utf-8 

"""
找出一颗二叉搜索树第K大的值
思路可以用中序遍历，第K个输出的结果就是第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.

Example 1:

Input: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
Output: 1
Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
Output: 3
@file: 230_KthSmallestElementinaBST.py
@time: 2019/5/14 18:04
"""

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        stack = []
        pos = root
        count = 0
        while pos or stack:
            if pos:
                stack.append(pos)
                pos = pos.left
            else:
                pos = stack.pop()
                count += 1
                if count == k:
                    return pos.val
                pos = pos.right