230. 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:

1. Try to utilize the property of a BST.
2. What if you could modify the BST node's structure?
3. The optimal runtime complexity is O(height of BST).

分析

在二叉检索树中查找第k小的元素，由于二叉检索树的左子树是较小元素存放位置，故使用中根遍历，每遍历一个元素k值减1，当k为0时即找到第k小的元素。

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if(root==NULL||k<1)
           return 0;
        stack<TreeNode*> s;//使用迭代法遍历二叉树 
        s.push(root);
        TreeNode* r=root->left;//指向左子树
        while(r||!s.empty())//设置r作为循环条件防止在[1,null,2]中出现1从栈中弹出，r指向2，但是在未压入栈前栈为空，无法进入循环的情况
        {
            while(r)//沿左子树向下查找
            {
                s.push(r);
                r=r->left;
            }
            r=s.top();//遍历根节点
            cout<<r->val<<endl;
            s.pop();
            k--;
            if(k==0)
                return r->val;
            r=r->right;//遍历右子树
        }
        return 1;
    }
};