leetcode 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.
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
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?
中文：给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

---

方法1：中序遍历（左->根->右）二叉搜索树再push_back进数组，即可得到一个从小到大排序的数组，在找第k个元素即可

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

方法二：利用二叉搜索树的性质，用分治算法。首先计算出左子树的结点个数总和cnt(定义count函数 count(root->left)，如果k小于等于左子树结点总和cnt，说明第k小的元素在左子树中，直接对左子结点调用递归即可。如果k大于cnt+1，说明目标值在右子树中，对右子结点调用递归函数，注意此时的k应为k-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);
        }
        return root->val;
    }
    int count(TreeNode* node) {
        if (!node) return 0;
        return 1 + count(node->left) + count(node->right);
    }
};