A Simple But Complete Implementation of Binary Search Tree

Binary Search Tree（BST）is a simple but complicated data structure. It can be modified dramatically to be a AVL or Black Red Tree in order to keep its performance and avoid being degraded to a single linked list. The following code is my implementation of a simple but complete binary search tree. It has common functionalities, like add, delete and search. Besides, I also implemented its iterator which made use of in-order traversal strategy to keep the result sorted in an ascending order.

Ok, It's time to share my code (Note: the following class is a generic class):

public class BST<E extends Comparable<? super E>> implements Iterable<E> {
	
	private TreeNode<E> root;   // the root of the tree
	
	// this is the data structure for each node in the tree
	private static class TreeNode<E> {
		E val;
		TreeNode<E> left;
		TreeNode<E> right;
		TreeNode(E x) { val = x; }
	}
	
	/**
	 * An iterator that iterates through a tree using in-order tree traversal
	 * allowing a sorted sequence.
	 *
	 */
	private static class BSTIter<E> implements Iterator<E> {
		TreeNode<E> current;
		Stack<TreeNode<E>> stack = new Stack<>();
		
		BSTIter(TreeNode<E> root) {
			this.current = root;
		}

		@Override
		public boolean hasNext() {
			return (!stack.empty() || current != null);
		}

		@Override
		public E next() {
			while (current != null) {
				stack.push(current);
				current = current.left;
			}
			
			current = stack.pop();
			TreeNode<E> node = current;
			current = current.right;
			
			return node.val;
		}
		
		@Override
		public void remove() {
			throw new UnsupportedOperationException("remove");
		}
	}
	
	public BST() {
		
	}
	
	public BST(E rootVal) {
		root = new TreeNode<E>(rootVal);
	}

	@Override
	public Iterator<E> iterator() {
		return new BSTIter<E>(root);
	}
	
	/**
	 * Add an element into the binary search tree
	 * @param element - the element to be added into the binary search tree
	 * @return true, if the given element is non-null
	 *         false, if the given element is null
	 */
	public boolean add(E element) {
		if (element == null) {
			return false;
		}
		
		if (root == null) {
			root = new TreeNode<E>(element);
		} else {
			TreeNode<E> parentNode = null;   // in adding the elements in BST, we must record the parent node for each loop.
			TreeNode<E> currentNode = root;
			
			while (currentNode != null) {
				parentNode = currentNode;
				if (element.compareTo(currentNode.val) <= 0) {
					currentNode = currentNode.left;
				} else {
					currentNode = currentNode.right;
				}
			}
			
			if (element.compareTo(parentNode.val) <= 0) {
				parentNode.left = new TreeNode<E>(element);
			} else {
				parentNode.right = new TreeNode<E>(element);
			}
			
		}
		
		return true;
	}
	
	/**
	 * Remove the given element from the binary search tree at the first time found
	 * @param element - the element to be removed from the binary search tree
	 * @return - the element removed from the binary search tree, if succeed
	 * 		     or null, if this element is not found in the binary search tree
	 */
	public E removeFirst(E element) {
		if (element == null) {
			return null;
		}
		
		TreeNode<E> parentNode = null;    // in removing the elements in BST, we must record the parent node for each loop.
		TreeNode<E> currentNode = root;
		
		while (currentNode != null) {
			if (element.compareTo(currentNode.val) < 0) {   // element must appear on the left side of the current node
				parentNode = currentNode;
				currentNode = currentNode.left;
			} else if (element.compareTo(currentNode.val) > 0) {  // element must appear on the right side of the current node
				parentNode = currentNode;
				currentNode = currentNode.right;
			} else {   // found the target element!!!
				// discuss three cases in which different deleting strategies are applied
				// 1. this removed node is the leaf node; 2. this removed node has just one child tree (left or right); 3. this removed node has two child trees
				if (currentNode.left == null && currentNode.right == null) {   // 1. this removed node is the leaf node
					if (parentNode == null) {
						// this node is root node
						root = null;
						return currentNode.val;
					} else {
						// this node is actual leaf node
						if (parentNode.left == currentNode) {  // to determine which side of parent node the current node belongs to
							parentNode.left = null;
							return currentNode.val;
						} else {
							parentNode.right = null;
							return currentNode.val;
						}
					}
				} else if (currentNode.left != null && currentNode.right == null) {   // 2. this removed node has just left child tree
					if (parentNode == null) {   // this node is root node
						root = currentNode.left;
						currentNode.left = null;   // avoid obsolete objects
						return currentNode.val;
					} else {
						if (parentNode.left == currentNode) {
							parentNode.left = currentNode.left;
							currentNode.left = null;
							return currentNode.val;
						} else {
							parentNode.right = currentNode.left;
							currentNode.left = null;
							return currentNode.val;
						}
					}
				} else if (currentNode.left == null && currentNode.right != null) {   // 2. this removed node has just right child tree
					if (parentNode == null) {   // this node is root node
						root = currentNode.right;
						currentNode.right = null;
						return currentNode.val;
					} else {
						if (parentNode.left == currentNode) {
							parentNode.left = currentNode.right;
							currentNode.right = null;
							return currentNode.val;
						} else {
							parentNode.right = currentNode.right;
							currentNode.right = null;
							return currentNode.val;
						}
					}
				} else {  // 3. this removed node has both left and right child trees
					/* steps: find the maximum value from this removed node's left child tree
					 * 		  replace this removed node's value with this maximum
					 *        remove the node who supply this maximum value (Note: this supplier node has only left child tree because of the property of bst, so remove this supplier is just like removing a node who has only one child tree)  
					 *  */
					TreeNode<E> leftChild = currentNode.left;
					TreeNode<E> leftChildParent = currentNode;
					
					while (leftChild.right != null) {
						leftChildParent = leftChild;
						leftChild = leftChild.right;
					}
					
					E returnVal = currentNode.val;
					currentNode.val = leftChild.val;
					
					if (leftChildParent == root) {   // do not forget to consider "<span style="font-family: Arial, Helvetica, sans-serif;">leftChildParent == root"</span>
						leftChildParent.left = leftChild.left;
						leftChild.left = null;
					} else {
						leftChildParent.right = leftChild.left;
						leftChild.left = null;
					}
					
					return returnVal;
				}
			}
		}
		
		return null;
	}

	/**
	 * Remove all the elements equal to the given elements from the bst
	 * @param element - the element to be removed from the binary search tree
	 * @return true, if remove at least one element
	 * 		   false, if no element removed
	 */
	public boolean removeAll(E element) {
		boolean found = false;
		
		if (removeFirst(element) != null) {
			while (true) {  // remove all the elements
				if (removeFirst(element) == null) {
					break;
				}
			}
			
			found = true;
		}
		
		return found;
	}
	
	/**
	 * To find a element equal to the given element
	 * @param element - the element to be found
	 * @return true, if found
	 * 	       otherwise, false
	 */
	public boolean findFirst(E element) {
		if (element == null) {
			return false;
		}
		
		TreeNode<E> node = root;
		while (node != null) {
			if (element.compareTo(node.val) < 0) {
				node = node.left;
			} else if (element.compareTo(node.val) > 0) {
				node = node.right;
			} else {
				return true;
			}
		}
		
		return false;
	}
}