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.

那麼這道題給的提示是讓我們用BST的性質來解題，最重要的性質是就是左<根<右，那麼如果用中序遍歷所有的節點就會得到一個有序陣列。所以解題的關鍵還是中序遍歷啊。關於二叉樹的中序遍歷可以參見我之前的部落格

解法一：

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

當然，此題我們也可以用遞迴來解，還是利用中序遍歷來解，程式碼如下：

解法二：

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

再來看一種分治法的思路，由於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);
        }
        return root->val;
    }
    int count(TreeNode* node) {
        if (!node) return 0;
        return 1 + count(node->left) + count(node->right);
    }
};

這道題的Follow up中說假設該BST被修改的很頻繁，而且查詢第k小元素的操作也很頻繁，問我們如何優化。其實最好的方法還是像上面的解法那樣利用分治法來快速定位目標所在的位置，但是每個遞迴都遍歷左子樹所有結點來計算個數的操作並不高效，所以我們應該修改原樹結點的結構，使其儲存包括當前結點和其左右子樹所有結點的個數，這樣我們使用的時候就可以快速得到任何左子樹結點總數來幫我們快速定位目標值了。定義了新結點結構體，然後就要生成新樹，還是用遞迴的方法生成新樹，注意生成的結點的count值要累加其左右子結點的count值。然後在求第k小元素的函式中，我們先生成新的樹，然後呼叫遞迴函式。在遞迴函式中，不能直接訪問左子結點的count值，因為左子節結點不一定存在，所以我們先判斷，如果左子結點存在的話，那麼跟上面解法的操作相同。如果不存在的話，當此時k為1的時候，直接返回當前結點值，否則就對右子結點呼叫遞迴函式，k自減1，參見程式碼如下：

解法四：

// Follow up
class Solution {
public:
    struct MyTreeNode {
        int val;
        int count;
        MyTreeNode *left;
        MyTreeNode *right;
        MyTreeNode(int x) : val(x), count(1), left(NULL), right(NULL) {}
    };
    
    MyTreeNode* build(TreeNode* root) {
        if (!root) return NULL;
        MyTreeNode *node = new MyTreeNode(root->val);
        node->left = build(root->left);
        node->right = build(root->right);
        if (node->left) node->count += node->left->count;
        if (node->right) node->count += node->right->count;
        return node;
    }
    
    int kthSmallest(TreeNode* root, int k) {
        MyTreeNode *node = build(root);
        return helper(node, k);
    }
    
    int helper(MyTreeNode* node, int k) {
        if (node->left) {
            int cnt = node->left->count;
            if (k <= cnt) {
                return helper(node->left, k);
            } else if (k > cnt + 1) {
                return helper(node->right, k - 1 - cnt);
            }
            return node->val;
        } else {
            if (k == 1) return node->val;
            return helper(node->right, k - 1);
        }
    }
};

類似題目：

參考資料：