-
Notifications
You must be signed in to change notification settings - Fork 0
/
BSTIterator.java
106 lines (96 loc) · 2.81 KB
/
BSTIterator.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
package binarysearchtreeiterator;
import datastructures.TreeNode;
import java.util.ArrayDeque;
import java.util.Stack;
/**
* Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.
* <p>
* Calling next() will return the next smallest number in the BST.
* <p>
* <p>
* <p>
* Example:
* <p>
* <p>
* <p>
* BSTIterator iterator = new BSTIterator(root);
* iterator.next(); // return 3
* iterator.next(); // return 7
* iterator.hasNext(); // return true
* iterator.next(); // return 9
* iterator.hasNext(); // return true
* iterator.next(); // return 15
* iterator.hasNext(); // return true
* iterator.next(); // return 20
* iterator.hasNext(); // return false
* <p>
* <p>
* Note:
* <p>
* next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
* You may assume that next() call will always be valid, that is, there will be at least a next smallest number in the BST when * next() is called.
*/
/**
* Your BSTIterator object will be instantiated and called as such:
* BSTIterator obj = new BSTIterator(root);
* int param_1 = obj.next();
* boolean param_2 = obj.hasNext();
*/
class BSTIterator {
ArrayDeque<TreeNode> nodes;
Stack<Integer> minNodes;
Stack<Integer> temp;
BSTIterator(TreeNode root) {
nodes = new ArrayDeque<>();
minNodes = new Stack<>();
temp = new Stack<>();
constructTree(root);
}
private void constructTree(TreeNode root) {
if (root == null) {
return;
}
nodes.add(root);
int count = 0;
TreeNode curr, left, right;
while (!nodes.isEmpty()) {
curr = nodes.poll();
// add curr on stack
if (minNodes.isEmpty()) {
minNodes.push(curr.val);
} else if (curr.val < minNodes.peek()) {
minNodes.push(curr.val);
} else {
while (!minNodes.isEmpty() && curr.val > minNodes.peek()) {
temp.push(minNodes.pop());
count++;
}
minNodes.push(curr.val);
for (int i = 0; i < count; i++) {
minNodes.push(temp.pop());
}
count = 0;
}
left = curr.left;
right = curr.right;
if (left != null) {
nodes.add(left);
}
if (right != null) {
nodes.add(right);
}
}
}
/**
* @return the next smallest number
*/
int next() {
return minNodes.isEmpty() ? -1 : minNodes.pop();
}
/**
* @return whether we have a next smallest number
*/
boolean hasNext() {
return !minNodes.isEmpty();
}
}