题目:
Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.
Calling next() will return the next smallest number in the BST.
思路:
用cur指向下一个节点。用my_stack(栈)只保存左子树的节点,目的是为了cur指针能够回退(因为是中序遍历所以需要回退)。如果下一个节点不存在,则cur会指向nullptr;否则将指向下一个节点。
代码实现:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode() : val(0), left(nullptr), right(nullptr) {}
* TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
* TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
* };
*/
class BSTIterator {
public:
stack<TreeNode*> my_stack;
TreeNode *cur;
BSTIterator(TreeNode* root) {
if (root == nullptr){
cur = nullptr;
}else{
while (root->left){
my_stack.push(root);
root = root->left;
}
cur = root;
}
}
/** @return the next smallest number */
int next() {
int ret = cur->val;
if (cur->right){
cur = cur->right;
while (cur->left){
my_stack.push(cur);
cur = cur->left;
}
}else{
if (my_stack.empty()){
cur = nullptr;
}else{
cur = my_stack.top();
my_stack.pop();
}
}
return ret;
}
/** @return whether we have a next smallest number */
bool hasNext() {
return cur != nullptr;
}
};
/**
* Your BSTIterator object will be instantiated and called as such:
* BSTIterator* obj = new BSTIterator(root);
* int param_1 = obj->next();
* bool param_2 = obj->hasNext();
*/