173. 二叉搜索树迭代器

173. 二叉搜索树迭代器

题目描述

解题思路

1、返回下一个最小数，BST中序遍历是一个升序序列。要求空间复杂度O(h)，所以需要自定义栈模拟BST的中序遍历,栈顶为遍历到当下next的最小值节点
2、当有右节点时，需要借助使用初始化入栈时的函数进行同样的操作

代码实现

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class BSTIterator:
    # 返回下一个最小数，BST中序遍历是一个升序序列。空间复杂度O(h)，所以需要自定义栈模拟BST的中序遍历
    # 栈顶为遍历到当下next的最小值节点
    def __init__(self, root: TreeNode):
        self.stack = []
        self.helper(root)
    
    def helper(self, root):
        tmp = root
        while tmp:
            self.stack.append(tmp)
            tmp = tmp.left
            

    def next(self) -> int:
        if self.hasNext:
            node = self.stack.pop()
            if node.right:
                self.helper(node.right)
            return node.val
        return -1


    def hasNext(self) -> bool:
        return len(self.stack) > 0


# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()

参考文献

1、官方题解