Kth Smallest Element in BST

LeetCode #230

Description:

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?

Example:

Note

Idea:

Inorder traverse. visit node的时候count++，如果count==k，就说明找到了。

function如果void return，比较容易写。如果return TreeNode， 递归时候注意处理各种return的情况。

Time O(n), 因为要一个一个找， 找不到就traverse全部。

Code:

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int count=0;
        TreeNode *result =inorderBST(root, k, count);

        if(result == NULL) 
            return -1;
        else
            return result->val;
    }

    TreeNode* inorderBST(TreeNode* root, int k, int& count){
        if(root == NULL) return NULL;

        TreeNode *left = inorderBST(root->left, k, count); 

        // Found the element from left subtree, return!
        if(left != NULL)
            return left;

        count++;
        if(count == k) 
            return root; // found, return

        // Not in left subtree, not in current root, so Return right subtree
        return inorderBST(root->right, k, count);
    }
};