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?

方法1: recursion

思路：

利用BST的性质，递归询问left child有多少个节点，如果大于k，向左，如果小于k，向右。但是这个方法是自上而下，大量重复统计子树的count。

易错点

1. 如果递归时候把k算对了，是不会出现在第一层的递归中出现nullptr错误的，k - count - 1这里特别容易错。

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

方法2: dfs （inorder）

思路：

整体的过程是一个left -> root -> right的顺序：left——先要知道左子树一共有多少节点，root——如果左子树刚好是k-1个，当前节点可以写入global，right——继续检查右子数有多少，最后返回是left+right+1。仅仅是在inorder的过程中截取了一个kth变量。

Complexity

Time complexity: O(n)
Space complexity: O(h)

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

方法3：iterative（inorder）

思路：

修改inorder的过程，中间加入计数，找到了直接返回。

class Solution1 {
public:
    int kthSmallest(TreeNode* root, int k) {
        int result;
        stack<TreeNode*> myStack;
        
        
        while(k > 0){
            if (root){
                myStack.push(root);
                root = root->left;
            }
            else{
                //while(!myStack.empty()){
                if (myStack.empty()) break;
                root = myStack.top();
                myStack.pop();
                result = root->val;
                k--;
                // if (root  nullptr){
                    root = root->right;
                //}
           // }
                
            }
            
            
        }
        return result;
    }
};

方法4: augment the TreeNode

思路：

follow up中问，如果经常增加删除，怎么优化。那每次query都记一遍数就不太有效率了，要把这个信息直接写入node，每次从子树删除，所有途径的节点都要修改count。

// 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);
        }
    }
};