Sunday, February 24, 2019

LeetCode 173. 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.
'''
# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

#Runtime: 88 ms, Memory Usage: 20 MB
#next() can use pop(0)
class BSTIterator1(object):

    def __init__(self, root):
        """
        :type root: TreeNode
        """
        self.i = -1
        self.vals = []
       
        def add_val(root):
            if root:
                self.vals.append(root.val)
                add_val(root.left)
                add_val(root.right)
       
        add_val(root)
        self.vals.sort()

    def next(self):
        """
        @return the next smallest number
        :rtype: int
        """
        self.i += 1
        if self.i < len(self.vals):
            return self.vals[self.i]
        else:
            return None
       

    def hasNext(self):
        """
        @return whether we have a next smallest number
        :rtype: bool
        """
        return self.i < len(self.vals)-1

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

No comments:

Post a Comment