题目
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()都完成一次中序遍历中的一个循环。其中,还需要注意,要把已经获取到的节点剪枝。
Python代码
# 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 = []
self.current = root
def next(self) -> int:
"""
@return the next smallest number
"""
while self.current:
self.stack.append(self.current)
self.current = self.current.left
self.current = self.stack.pop()
# 剪枝步骤
if self.current and self.current.right:
self.current.left = None
val = self.current.val
self.current = self.current.right
return val
def hasNext(self) -> bool:
"""
@return whether we have a next smallest number
"""
if self.current or len(self.stack) != 0:
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()