From b339299d158b4399796cdcbb5f0e28ca986fa192 Mon Sep 17 00:00:00 2001
From: Abhinav Sarkar <abhinav@abhinavsarkar.net>
Date: Mon, 1 Jul 2019 12:51:57 +0530
Subject: [PATCH] Adds BST inorder and preorder iterators

---
 .../algorist/BinarySearchTree.java            | 156 ++++++++++++++++--
 1 file changed, 142 insertions(+), 14 deletions(-)

diff --git a/src/main/java/net/abhinavsarkar/algorist/BinarySearchTree.java b/src/main/java/net/abhinavsarkar/algorist/BinarySearchTree.java
index 61a3e01..7df9ca0 100644
--- a/src/main/java/net/abhinavsarkar/algorist/BinarySearchTree.java
+++ b/src/main/java/net/abhinavsarkar/algorist/BinarySearchTree.java
@@ -1,8 +1,11 @@
 package net.abhinavsarkar.algorist;
 
+import java.util.Iterator;
+import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.Stack;
 
-public class BinarySearchTree<K extends Comparable<K>,V>
+public class BinarySearchTree<K extends Comparable<K>,V> implements Iterable<BinarySearchTree.Entry<K,V>>
 {
     private Node<K,V> root = EmptyNode.instance();
 
@@ -38,6 +41,12 @@ public class BinarySearchTree<K extends Comparable<K>,V>
         return this.root.toStringBuilder(0).toString();
     }
 
+    @Override
+    public Iterator<Entry<K, V>> iterator()
+    {
+        return new PreOrderIterator<>(this.root);
+    }
+
     private interface Node<K extends Comparable<K>, V>
     {
         Node<K,V> insert(K key, V val);
@@ -48,9 +57,16 @@ public class BinarySearchTree<K extends Comparable<K>,V>
         Optional<K> successor(K key);
         Optional<K> predecessor(K key);
         boolean isEmpty();
-        ValueNode<K,V> toValueNode();
         StringBuilder toStringBuilder(int level);
         boolean checkIfBST();
+
+        default ValueNode<K,V> toValueNode() {
+            if (this.isEmpty()) {
+                throw new UnsupportedOperationException();
+            } else {
+                return (ValueNode<K, V>) this;
+            }
+        }
     }
 
     private static class ValueNode<K extends Comparable<K>,V> implements Node<K, V>
@@ -78,12 +94,6 @@ public class BinarySearchTree<K extends Comparable<K>,V>
             return false;
         }
 
-        @Override
-        public ValueNode<K, V> toValueNode()
-        {
-            return this;
-        }
-
         @Override
         public Node<K, V> insert(K key, V val)
         {
@@ -250,12 +260,6 @@ public class BinarySearchTree<K extends Comparable<K>,V>
             return true;
         }
 
-        @Override
-        public ValueNode<K, V> toValueNode()
-        {
-            throw new UnsupportedOperationException();
-        }
-
         @Override
         public Optional<V> get(K key)
         {
@@ -315,6 +319,127 @@ public class BinarySearchTree<K extends Comparable<K>,V>
         return sb.append("├ ");
     }
 
+    public static class Entry<K, V> {
+        private K key;
+        private V val;
+
+        public Entry(K key, V val) {
+            this.key = key;
+            this.val = val;
+        }
+
+        public K getKey()
+        {
+            return key;
+        }
+
+        public V getVal()
+        {
+            return val;
+        }
+
+        public String toString() {
+            return "<" + key + "," + val + ">";
+        }
+    }
+
+    private static class InOrderIterator<K extends Comparable<K>, V> implements Iterator<Entry<K, V>>
+    {
+        private final Stack<Node<K, V>> stack = new Stack<>();
+
+        public InOrderIterator(Node<K, V> node)
+        {
+            pushToStack(node);
+        }
+
+        @Override
+        public boolean hasNext()
+        {
+            return !(this.stack.isEmpty() || this.stack.peek().isEmpty());
+        }
+
+        /*
+        inorder(a):
+            if a.left != null:
+                yield inorder(a.left)
+            yield a.val
+            if a.right != null:
+                yield inorder(a.right)
+         */
+        @Override
+        public Entry<K, V> next()
+        {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
+            }
+
+            ValueNode<K, V> vNode = this.stack.pop().toValueNode();
+            pushToStack(vNode.right);
+
+            return new Entry<>(vNode.key, vNode.val);
+        }
+
+        private void pushToStack(Node<K, V> node)
+        {
+            Node<K, V> current = node;
+            while (!current.isEmpty())
+            {
+                this.stack.push(current);
+                current = current.toValueNode().left;
+            }
+        }
+    }
+
+    private static class PreOrderIterator<K extends Comparable<K>, V> implements Iterator<Entry<K, V>>
+    {
+        private final Stack<Node<K, V>> stack = new Stack<>();
+        private Node<K, V> current;
+
+        public PreOrderIterator(Node<K, V> node)
+        {
+            current = node;
+        }
+
+        @Override
+        public boolean hasNext()
+        {
+            return !this.current.isEmpty();
+        }
+
+        /*
+        preorder(a):
+            yield a.val
+            if a.left != null:
+                yield preorder(a.left)
+            if a.right != null:
+                yield preorder(a.right)
+         */
+        @Override
+        public Entry<K, V> next()
+        {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
+            }
+
+            ValueNode<K, V> vNode = current.toValueNode();
+            if (!vNode.left.isEmpty()) {
+                stack.push(current);
+                current = vNode.left;
+            } else if (!vNode.right.isEmpty()) {
+                current = vNode.right;
+            } else if (stack.isEmpty()){
+                current = EmptyNode.instance();
+            } else {
+                do {
+                    Node<K, V> node = stack.pop();
+                    current = node.toValueNode().right;
+                } while (current.isEmpty());
+            }
+
+            return new Entry<>(vNode.key, vNode.val);
+        }
+
+    }
 
     public static void main(String[] args)
     {
@@ -371,5 +496,8 @@ public class BinarySearchTree<K extends Comparable<K>,V>
         System.out.println(bst.predecessor("z"));
         System.out.println(bst.predecessor("q"));
 
+        for (Entry<String,String> entry: bst) {
+            System.out.println(entry.key);
+        }
     }
 }