Description
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.
Example:
BSTIterator iterator = new BSTIterator(root);
iterator.next(); // return 3
iterator.next(); // return 7
iterator.hasNext(); // return true
iterator.next(); // return 9
iterator.hasNext(); // return true
iterator.next(); // return 15
iterator.hasNext(); // return true
iterator.next(); // return 20
iterator.hasNext(); // return false
Note:
next()
andhasNext()
should run in average O(1) time and uses O(h) memory, where h is the height of the tree.- You may assume that
next()
call will always be valid, that is, there will be at least a next smallest number in the BST whennext()
is called.
Answer:
第一反应还是用栈来做。先递归把左子树全部压进去,之后每调用一次next()
,就出栈,同时把出栈的所有左路径上的节点压入。
而hasNext()
则是返回栈是否为空。
下面是代码:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class BSTIterator {
public:
void push(TreeNode* root) {
while (root) {
sta.push(root);
root = root -> left;
}
}
BSTIterator(TreeNode* root) {
push(root);
}
/** @return the next smallest number */
int next() {
TreeNode* temp = sta.top();
sta.pop();
if (temp -> right)
push(temp -> right);
return temp -> val;
}
/** @return whether we have a next smallest number */
bool hasNext() {
return !sta.empty();
}
private:
stack<TreeNode*> sta;
int size;
int pos;
};
/**
* 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();
*/
上述next()
的空间复杂度是O(h),而通过摊还分析,我们也可以计算出其时间复杂度是O(1),所以是符合题目要求的。
不过之后想了一想,既然题目至少要求next()
和hasNext()
的时间复杂度和空间复杂度,那么构造函数的复杂度就是没有限制的。所以我们可以直接在构造函数中进行中序遍历,这样就可以保证next()
和hasNext()
的空间和时间都是常数复杂度了。
下面是代码:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class BSTIterator {
public:
vector<int> arr;
int pos;
void compute(TreeNode* root) {
if (!root)
return;
compute(root -> left);
arr.push_back(root -> val);
compute(root -> right);
}
BSTIterator(TreeNode* root) {
compute(root);
pos = 0;
}
/** @return the next smallest number */
int next() {
return arr[pos++];
}
/** @return whether we have a next smallest number */
bool hasNext() {
return pos < arr.size();
}
};
/**
* 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();
*/