LeetCode - Binary Search Tree Iterator (Java)
Tags: binary search tree, in order traversal, iterator, LeetCode, TreeBinary Search Tree Iterator
- Difficulty: Medium
Implement 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.
Source: https://leetcode.com/problems/binary-search-tree-iterator/
In this post, we describe a stack based method to solve the Binary Search Tree Iterator problem in Java. This solution similar to the stack based in oder traversal of a binary search tree.
Analyis
Since “Calling next() will return the next smallest number in the BST”, it is actually an in order traverse of the binary search tree. It is easy to implement the in order binary search tree traversal algorithm using a stack. However, to build a iterator for a binary search tree, we need to figure out which part of the code should be put into the hasNext() method, and which part of the code should be put into the next() method.
Since we start from the root of tree, we can first push all the left children of root to the stack, now we know the top of the stack is the minimum value of the tree.
The code is like this:
|
1 2 3 4 5 6 7 8 9 |
public BSTIterator(TreeNode root) { stack = new Stack<TreeNode>(); while(root != null) { stack.push(root); root = root.left; } } |
In the hasNext() method, we simply check whether the stack is empty.
|
1 2 3 4 5 |
/** @return whether we have a next smallest number */ public boolean hasNext() { return !stack.isEmpty(); } |
In the next() method, we first pop up the stack to get the node, which is the current smallest node. Then we need to check whether this node has right child. If it does have right child, we need to put its right child and the right child’s left branch to the stack. After this process, the top of the stack holds the next smallest node.
|
1 2 3 4 5 6 7 8 9 10 |
public int next() { TreeNode n = stack.pop(); // n the node with the smallest value TreeNode r = n.right; while(r!=null){ stack.push(r); r = r.left; } return n.val; } |
Here is the full implementation of the Binary Search Tree Iterator in 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 |
public class BSTIterator { Stack<TreeNode> stack; public BSTIterator(TreeNode root) { stack = new Stack<TreeNode>(); while(root != null) { stack.push(root); root = root.left; } } /** @return whether we have a next smallest number */ public boolean hasNext() { return !stack.isEmpty(); } /** @return the next smallest number */ public int next() { TreeNode n = stack.pop(); TreeNode r = n.right; while(r!=null){ stack.push(r); r = r.left; } return n.val; } } |
