【python/leetcode/M】Binary Search Tree Iterator

Keywords: Programming

subject

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.
Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
Credits:
Special thanks to @ts for adding this problem and creating all test cases.

Implementation ideas

To find the smallest node in the binary search tree, we should recursively traverse the left node from the root node until the traversed node has no left node, then the node is the smallest node in the binary tree. Now that there is the smallest node in the binary search tree that has not been visited, where is the smallest node in the remaining tree that has not been visited after visiting the node? If the node has a right subtree, it is in its right subtree (returning to the smallest element of finding a binary search tree, but the binary search tree becomes smaller); if there is no right subtree, it is its parent.

In order to locate the parent node quickly, we can temporarily store the traversal path with the stack. When we do next(), we pop up the top element of the stack and access it. If it has a right subtree, we traverse its right subtree. If there is no right subtree, we will visit the parent node of the current node the next time we exit the stack.

Implementation code

# Definition for a  binary tree node
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class BSTIterator(object):
    def __init__(self, root):
        """
        :type root: TreeNode
        """
        self.stack = []
        self._pushLeft(root)
        

    def hasNext(self):
        """
        :rtype: bool
        """
        return self.stack
        

    def next(self):
        """
        :rtype: int
        """
        node = self.stack.pop()
        self._pushLeft(node.right)
        return node.val
    
    
    def _pushLeft(self,node):
        while node:
            self.stack.append(node)
            node = node.left
            
# Your BSTIterator will be called like this:
# i, v = BSTIterator(root), []
# while i.hasNext(): v.append(i.next())

Posted by ericm on Wed, 23 Jan 2019 03:18:13 -0800