binary search tree iterator

🏠

Tushar Solution

Credit: Tushar

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

 1 class TreeNode:
 2     def __init__(self, x, left=None, right=None):
 3         self.val = x
 4         self.left = left
 5         self.right = right
 6 
 7 class 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 
28 obj = BSTIterator(TreeNode(7, TreeNode(3), TreeNode(15, TreeNode(9), TreeNode(20))))
29 
30 for v in obj:
31     print(v)

Output:

1 sherlock ran 33 lines of Python 3:
2 
3 3
4 7
5 9
6 15
7 20
8 
9 >>>                               (finished in 1.21s):            

Shyal Solution

 1 class TreeNode:
 2     def __init__(self, val=0, left=None, right=None):
 3         self.val = val
 4         self.left = left
 5         self.right = right
 6 
 7 class 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