使用栈来记录可能的路径,栈顶一直是下一个元素。
class BSTIterator {
public:
stack<TreeNode *> path;
BSTIterator(TreeNode *root) {
path = stack<TreeNode *>();
TreeNode *node = root;
while (node != NULL) {
path.push(node);
node = node->left;
}
}
/** @return whether we have a next smallest number */
bool hasNext() {
return !path.empty();
}
/** @return the next smallest number */
int next() {
TreeNode *node = path.top();
path.pop();
int ret = node->val;
if (node->right != NULL) {
node = node->right;
path.push(node);
while (node->left != NULL) {
node = node->left;
path.push(node);
}
}
return ret;
}
};
Python3
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class BSTIterator:
def __init__(self, root: TreeNode):
self.stack = [] # append, pop
self.node = root
while self.node is not None:
self.stack.append(self.node)
self.node = self.node.left
def next(self) -> int:
"""
@return the next smallest number
"""
if len(self.stack) > 0:
self.node = self.stack.pop()
result = self.node.val
self.node = self.node.right
while self.node:
self.stack.append(self.node)
self.node = self.node.left
return result
def hasNext(self) -> bool:
"""
@return whether we have a next smallest number
"""
if len(self.stack):
return True
return False
# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()