Kth Smallest Element in a BST

题目描述：

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?

Hint:

Try to utilize the property of a BST.

What if you could modify the BST node's structure?

The optimal runtime complexity is O(height of BST).

题目大意：

给定一棵二叉搜索树（BST），编写一个函数kthSmallest找出其中第k小的元素。

注意：

你可以假设k总是有效的， 1 ≤ k ≤ BST的元素总数。

进一步思考：

如果BST的修改（插入/删除）操作十分频繁，并且需要频繁地找出第k小的元素，应该怎样优化kthSmallest函数？

提示：

尝试利用BST的属性。

如果你可以修改BST节点的结构时，应该怎样做？

最优时间复杂度应该是O(BST的高度)。

解题思路：

BST具有如下性质：

左子树中所有元素的值均小于根节点的值

右子树中所有元素的值均大于根节点的值

因此采用中序遍历（左 -> 根 -> 右）即可以递增顺序访问BST中的节点，从而得到第k小的元素，时间复杂度O(k)

Python代码：

# Definition for a binary tree node.
# class TreeNode:
#	 def __init__(self, x):
#		 self.val = x
#		 self.left = None
#		 self.right = None
class Solution:
  # @param {TreeNode} root
  # @param {integer} k
  # @return {integer}
  def kthSmallest(self, root, k):
    stack = []
    node = root
    while node:
      stack.append(node)
      node = node.left
    x = 1
    while stack and x <= k:
      node = stack.pop()
      x += 1
      right = node.right
      while right:
        stack.append(right)
        right = right.left
    return node.val

进一步思考：

如果BST节点TreeNode的属性可以扩展，则再添加一个属性leftCnt，记录左子树的节点个数

记当前节点为node

当node不为空时循环：

若k == node.leftCnt + 1：则返回node

否则，若k > node.leftCnt：则令k -= node.leftCnt + 1，令node = node.right

否则，node = node.left

上述算法时间复杂度为O(BST的高度)