binary search tree iterator

🏠

Tushar Solution

Credit: Tushar

https://leetcode.com/problems/binary-search-tree-iterator/

 1class TreeNode:
 2    def __init__(self, x, left=None, right=None):
 3        self.val = x
 4        self.left = left
 5        self.right = right
 6
 7class BSTIterator:
 8    def __init__(self, root: TreeNode):
 9        self.root = root
10        self.gen = self.dfs(self.root)
11
12    def dfs(self, n):
13        if n:
14            if n.left:
15                yield from self.dfs(n.left)
16            yield n.val
17            if n.right:
18                yield from self.dfs(n.right)
19
20    def __iter__(self):
21        return self
22
23    def __next__(self) -> int:
24        return next(self.gen)
25
26
27
28obj = BSTIterator(TreeNode(7, TreeNode(3), TreeNode(15, TreeNode(9), TreeNode(20))))
29
30for v in obj:
31    print(v)

Output:

1sherlock ran 33 lines of Python 3:
2
33
47
59
615
720
8
9>>>                               (finished in 1.21s):            

Shyal Solution

 1class TreeNode:
 2    def __init__(self, val=0, left=None, right=None):
 3        self.val = val
 4        self.left = left
 5        self.right = right
 6
 7class BSTIterator:
 8    def __init__(self, root: TreeNode):
 9        self.root = root
10        self.iter = self.dfs(self.root)
11        self.next_item = None
12
13    def dfs(self, n):
14        if n:
15            yield from self.dfs(n.left)
16            yield n.val
17            yield from self.dfs(n.right)
18        
19    def next(self) -> int:
20        """
21        Returns the next item if it's been cached,
22        else calls next.
23        """
24        if self.next_item is not None:
25            tmp, self.next_item = self.next_item, None
26            return tmp
27        return next(self.iter)
28
29    def hasNext(self) -> bool:
30        """
31        Returns whether we have a next item in the iterator.
32        Does so by actually calling next, and storing the result
33        in a cache.
34        """
35        if self.next_item is None:
36            self.next_item = next(self.iter, None)
37        return self.next_item is not None
38