LeetCode Binary Search Tree Iterator

LeetCode Binary Search Tree Iterator

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.

Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.


本题要求设计一个二叉搜索树的迭代器,迭代的输出当前树中最小值,并且要求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。

Leave a Reply

Your email address will not be published. Required fields are marked *