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.
本题要求设计一个二叉搜索树的迭代器,迭代的输出当前树中最小值,并且要求next和hasNext是O(1)时间复杂度和O(h)空间复杂度,h为树的高度。 因为二叉搜索树的中序遍历是递增有序的序列,所以如果没有O(h)空间复杂度限制,则可以提前保存中序遍历结果,然后实现对数组的迭代器。但是这样的空间复杂度就是O(sizeof(tree)),不符合题意。 立马想到树的中序遍历的非递归实现,我们可以借助一个堆栈,先只把树的根和最左边的数压栈,这样栈顶元素就是树的最下角元素,也就是最小值。每调用一次next,返回当前栈顶元素,同时把栈顶元素的右孩子及右孩子的所有最左边的节点压栈。hasNext就好办,直接返回堆栈是否为空。 这样堆栈中最多会保存树高个数的节点,所以是O(h)的空间复杂度。完整代码如下:
class BSTIterator {
private:
stack<TreeNode*> tree;
void putAllLeft(TreeNode* root, stack<TreeNode*>& tree)
{
while (root != NULL) {
tree.push(root);
root = root->left;
}
}
public:
BSTIterator(TreeNode* root) { putAllLeft(root, tree); } /** @return whether we have a next smallest number */
bool hasNext() { return !tree.empty(); } /** @return the next smallest number */
int next()
{
TreeNode* top = tree.top();
tree.pop();
putAllLeft(top->right, tree);
return top->val;
}
};
本代码提交AC,用时19MS。