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