[LeetCode]Binary Search Tree Iterator,答题报告
[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<TreeNode> mArrayDeque; public BSTIterator(TreeNode root) { mArrayDeque = new ArrayDeque<TreeNode>(); bTreeInorderTraverse(root); } private void bTreeInorderTraverse(TreeNode root) { TreeNode p = root; Stack<TreeNode> tmpStack = new Stack<TreeNode>(); 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<TreeNode> mStack; public BSTIterator(TreeNode root) { mStack = new Stack<TreeNode>(); 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; } } }