[LeetCode] 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).

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

分析：

方法一：中序遍历+直接去第k个；

方法二：使用中序遍历框架，当检测到第k个时，直接返回，空间 复杂度节省一些。

中序遍历的code

class Solution {
public:
    vector<int> preorderTraversal(TreeNode *root) {
        
        vector<int> result;
        stack<TreeNode*> stack;
        
        TreeNode *p = root;
        
        while( NULL != p || !stack.empty())
        {
            while(NULL != p)
            {
                result.push_back(p->val);
                
                stack.push(p);
                p = p->left;
            }
            
            if(!stack.empty())
            {
                p= stack.top();
                stack.pop();
                
                p=p->right;
            }
        }
        
        return result;
    }
    
};

改进的使用中序遍历来解决此问题的code：

class Solution {
    public:
        int kthSmallest(TreeNode* root, int k) {
    
            stack<TreeNode*> st; 
            TreeNode* p = root;

            int cnt = 0;

            while(p != NULL || !st.empty())
            {   
                while(p)
                {   
                    st.push(p);
                    p = p->left;
                }   

                if(!st.empty())
                {   
                    p = st.top();
                    st.pop();

                    cnt ++; 
                    if(cnt == k)
                    {   
                        return p->val;
                    }   

                    p = p->right;

                }   
            }
            return -1;

        }
};