[LeetCode]Binary Search Tree Iterator,解題報告
題目
LeetCode題目如下:
mplement 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.
思路
其實LeetCode的題目並不難,很多第一眼沒有思路的題目畫個圖就清楚了。這道題也是如此,想實現二叉搜索樹的Iterator,每次找當前的最小值。畫個圖就知道,這個題目就是考察二叉樹的中序遍歷而已。再不考慮空間復雜度的情況下,我們可以很簡單的寫出二叉搜索樹的AC代碼:
import java.util.ArrayDeque;
import java.util.Stack;
public class BSTIterator {
private ArrayDeque mArrayDeque;
public BSTIterator(TreeNode root) {
mArrayDeque = new ArrayDeque();
bTreeInorderTraverse(root);
}
private void bTreeInorderTraverse(TreeNode root) {
TreeNode p = root;
Stack tmpStack = new Stack();
while (p != null || ! tmpStack.empty()) {
if (p != null) {
tmpStack.push(p);
p = p.left;
} else {
p = tmpStack.pop();
mArrayDeque.add(p);
p = p.right;
}
}
}
public boolean hasNext() {
return ! mArrayDeque.isEmpty();
}
public int next() {
if (hasNext()) {
return mArrayDeque.poll().val;
} else {
return -1;
}
}
public static class TreeNode {
int val;
TreeNode left;
TreeNode right;
TreeNode(int x) {
val = x;
}
}
}
優化
題目中給出的空間復雜度為O(h),而我使用的空間復雜度為O(2n),不符合題目的要求。因此需要考慮如何修改代碼邏輯解決這個問題。思路還是使用二叉樹的中序遍歷,但是在空間有限的情況下(ps:只有O(h)),我們可以不在構造函數中完成所有的中序遍歷操作。思路如下:
申請一個只有h大小的棧。在構造函數中,將該樹root節點的所有左孩子入棧。在next()函數中,彈出棧頂節點,如果該節點有右孩子,將右孩子和該右孩子的所有左孩子入棧。
思路很簡單,說白了,就是將在中序遍歷算法分解,在構造函數和next方法中共同完成。AC代碼如下:
import java.util.Stack;
public class BSTIterator {
private Stack mStack;
public BSTIterator(TreeNode root) {
mStack = new Stack();
while (root != null) {
mStack.push(root);
root = root.left;
}
}
public boolean hasNext() {
return ! mStack.empty();
}
public int next() {
TreeNode p = mStack.pop();
int res = p.val;
if (p.right != null) {
TreeNode node = p.right;
while (node != null) {
mStack.push(node);
node = node.left;
}
}
return res;
}
public static class TreeNode {
int val;
TreeNode left;
TreeNode right;
TreeNode(int x) {
val = x;
}
}
}