Reimplemented remove operation with remove root element

Author: Swapnil Gaikwad

src/main/java/com/datastructures/Tree.java

@@ -33,90 +33,59 @@ public void add( E data ) {
         }
     }
 
-    public boolean find( E data ) {
-        return find( data, root );
+    public boolean contains( E data ) {
+        return contains( data, root );
     }
 
-    private boolean find( E data, Node root ) {
+    private boolean contains( E data, Node root ) {
         if ( root == null ) {
             return false;
         }
         if ( root.data == data ) {
             return true;
         }
         if ( data.compareTo( root.data ) < 0 ) {
-            return find( data, root.left );
+            return contains( data, root.left );
         } else {
-            return find( data, root.right );
+            return contains( data, root.right );
         }
     }
 
     public boolean remove( E data ) {
-        if ( root.data == data ) {
-            return false;
-        }
-        Node parent = findParent( root, data );
-        if ( parent == null ) {
-            return false;
-        }
-        Node current;
-        if ( parent.left != null && parent.left.data == data ) {
-            current = parent.left;
-        } else {
-            current = parent.right;
-        }
-        remove( parent, current );
-        return true;
+        final boolean isPresent = contains( data );
+        root = remove( root, data );
+        return isPresent;
     }
 
-    private void remove( Node parent, Node current ) {
-        if ( current.isLeaf() ) {
-            if ( parent.left == current ) {
-                parent.left = null;
-            } else {
-                parent.right = null;
-            }
-            return;
+    private Node remove( Node root, E data ) {
+        // If tree is not empty
+        if ( root == null ) {
+            return null;
         }
-        if ( current.left == null ) {
-            if ( parent.left == current ) {
-                parent.left = current.right;
-            } else {
-                parent.right = current.right;
+
+        if ( data.compareTo( root.data ) < 0 ) {
+            root.left = remove( root.left, data );
+        } else if ( data.compareTo( root.data ) > 0 ) {
+            root.right = remove( root.right, data );
+        } else {
+            if ( root.left == null ) {
+                return root.right;
             }
-            return;
-        }
-        if ( current.right == null ) {
-            if ( parent.left == current ) {
-                parent.left = current.left;
-            } else {
-                parent.right = current.left;
+            if ( root.right == null ) {
+                return root.left;
             }
-            return;
-        }
-        Node itr = current.right;
-        parent = current;
-        while ( itr.left != null ) {
-            parent = itr;
-            itr = itr.left;
+            //Element found
+            root.data = minValue( root.right );
+            root.right = remove( root.right, root.data );
         }
-        current.data = itr.data;
-        remove( parent, itr );
+        return root;
     }
 
-    private Node findParent( Node current, E data ) {
-        if ( current == null || current.data == data ) {
-            return null;
-        }
-        if ( current.left != null && current.left.data == data ) {
-            return current;
-        } else if ( current.right != null && current.right.data == data ) {
-            return current;
-        } else if ( data.compareTo( current.data ) < 0 ) {
-            return findParent( current.left, data );
-        } else {
-            return findParent( current.right, data );
+    private E minValue( Node root ) {
+        while ( root.left != null ) {
+            root = root.left;
         }
+        return root.data;
     }
 
     private class Node {

src/test/java/com/datastructures/TreeTest.java

@@ -16,8 +16,8 @@ public void add_node_in_tree() {
         tree.add( 2 );
         tree.add( 1 );
 
-        assertEquals( true, tree.find( 2 ) );
-        assertEquals( false, tree.find( 5 ) );
+        assertEquals( true, tree.contains( 2 ) );
+        assertEquals( false, tree.contains( 5 ) );
     }
 
     @Test
@@ -30,12 +30,12 @@ public void remove_node_in_tree() {
         tree.add( 1 );
         tree.add( 3 );
 
-        List<Integer> list = Arrays.asList( 8, 7, 2, 1, 3 );
+        List<Integer> list = Arrays.asList( 4, 8, 7, 2, 1, 3 );
         list.forEach( element -> {
             System.out.println( element );
-            assertEquals( true, tree.find( element ) );
+            assertEquals( true, tree.contains( element ) );
             assertEquals( true, tree.remove( element ) );
-            assertEquals( false, tree.find( element ) );
+            assertEquals( false, tree.contains( element ) );
         } );
     }
 }