[LeetCode] 230. Kth Smallest Element in a BST

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?

解析

找到二叉搜索树中第k小的点，中序遍历。

解法1：迭代

直接使用非递归的中序遍历方法，设置一个计数值，当i等于k时，输出结果。

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> s;
        if(root) s.push(root);
        TreeNode* p = root;
        int i = 0;
        while(p || !s.empty()){
            if(p){
                s.push(p);
                p = p->left;
            }
            else{
                p = s.top();
                s.pop();
                i ++;
                if(i == k)
                    return p->val;
                p = p->right;
            }
        }
    }
};

解法2：递归1

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        return kthSmallestDFS(root, k);
    }
    int kthSmallestDFS(TreeNode* root, int &k){
        if(!root) return -1;
        int left = kthSmallestDFS(root->left, k);
        if(k == 0) return left;
        if(--k == 0) return root->val;
        return kthSmallestDFS(root->right, k);
    }  
};

解法3：递归分治

由于BST的性质，我们可以快速定位出第k小的元素是在左子树还是右子树，我们首先计算出左子树的结点个数总和cnt，如果k小于等于左子树结点总和cnt，说明第k小的元素在左子树中，直接对左子结点调用递归即可。如果k大于cnt+1，说明目标值在右子树中，对右子结点调用递归函数，注意此时的k应为k-cnt-1，应为已经减少了cnt+1个结点。如果k正好等于cnt+1，说明当前结点即为所求，返回当前结点值即可。

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int cnt = count(root->left);
        if(k <= cnt)
            return kthSmallest(root->left, k);
        else if(k > cnt + 1)
            return kthSmallest(root->right, k-cnt-1);
        else return root->val;
    }
    int count(TreeNode* root){
        if(!root) return 0;
        else return 1 + count(root->left) + count(root->right);
    }  
};

参考

http://www.cnblogs.com/grandyang/p/4620012.html