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?
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
stack<TreeNode *> st;
TreeNode *p = root;
while(p || !st.empty())
{
while(p)
{
st.push(p);
p = p->left;
}
p = st.top();
if(--k == 0)
return p->val;
st.pop();
p = p->right;
}
}
};
const BinaryTreeNode* KthNodeCore(const BinaryTreeNode* pRoot, unsigned int& k)
{
const BinaryTreeNode* target = nullptr;
if(pRoot->m_pLeft != nullptr)
target = KthNodeCore(pRoot->m_pLeft, k);
if(target == nullptr)
{
if(k == 1)
target = pRoot;
k--;
}
if(target == nullptr && pRoot->m_pRight != nullptr)
target = KthNodeCore(pRoot->m_pRight, k);
return target;
}