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.

• 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:

  • Try to utilize the property of a BST.
  • What if you could modify the BST node's structure?
  • The optimal runtime complexity is O(height of BST).

Solution

1. use inorder traversal result
2. use the number of elements in left sub-tree

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

Solution, recursive inorder

class Solution {
public:
    void inorderTraverse(TreeNode* root, int k, vector<int>& inorder) {
        if ( root == NULL ) return;
        inorderTraverse(root->left, k, inorder);
        int n = inorder.size();
        if ( n == k ) return;
        inorder.push_back(root->val);
        if ( n == k-1 ) return;
        inorderTraverse(root->right, k, inorder);
        return;
    }
    int kthSmallest(TreeNode* root, int k) {
        vector<int> inorder;
        inorderTraverse(root, k, inorder);
        return inorder.back();
    }
};