需求
next() and hasNext() should run in average O(1) time and uses O(h) memory
算法
中序迭代遍历二叉树
时间复杂度:
- BSTIterator:
O(h)
- next: 平均
O(1)
最差O(h)
- hasNext:
O(1)
空间复杂度: O(h)
C++代码
class BSTIterator {
public:
BSTIterator(TreeNode *root) { initStack(root); }
int next() {
TreeNode *node = stk.top();
stk.pop();
initStack(node->right);
return node->val;
}
bool hasNext() const { return !stk.empty(); }
private:
inline void initStack(TreeNode *node) {
while (node) {
stk.push(node);
node = node->left;
}
}
stack<TreeNode *> stk;
};
Java代码
class BSTIterator {
public BSTIterator(TreeNode root) {
stk = new ArrayList<>();
initStack(root);
}
public int next() {
final TreeNode node = stk.get(stk.size() - 1);
stk.remove(stk.size() - 1);
initStack(node.right);
return node.val;
}
public boolean hasNext() {
return !stk.isEmpty();
}
private final void initStack(TreeNode node) {
while (node != null) {
stk.add(node);
node = node.left;
}
}
private List<TreeNode> stk;
}
Python3代码
class BSTIterator:
def __init__(self, root: TreeNode):
self.stk = []
self.__initStack(root)
def next(self) -> int:
node = self.stk.pop()
self.__initStack(node.right)
return node.val
def hasNext(self) -> bool:
return bool(self.stk)
def __initStack(self, node) -> None:
while node:
self.stk.append(node)
node = node.left