MyTreeSet.java

package week6.home_work;

import java.util.*;

/**
 * Created by nagornyyalek on 16.09.2016.
 */
//public class MyTreeSetInteger {
public class MyTreeSet<E> implements Set<E> {
    private TreeNode root;
    private int size;
    private Comparator<? super E> comparator;

    public MyTreeSet() {
    }

    public MyTreeSet(Comparator<? super E> comparator) {
        this.comparator = comparator;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public boolean contains(Object o) {
        return findNode(root, (E) o) != null;
    }

    private int canCompare(E element1, E element2) {
        return comparator != null ? comparator.compare(element1, element2) :
                ((Comparable<? super E>)element1).compareTo(element2);
    }

    private TreeNode findNode(TreeNode curr, E value) {

        if (curr == null) {
            return null;
        }

        int compare = canCompare(value, (E) curr.value);
       /* if (value < curr.value) {
            return findNode(curr.left, value);
        } else if (value > curr.value) {
            return findNode(curr.right, value);
        }*/

        if (compare < 0) {
            return findNode(curr.left, value);
        } else if (compare > 0) {
            return findNode(curr.right, value);
        }

        return curr;
    }

    @Override
    public Iterator<E> iterator() {

        return new MyTreeSetIterator();
    }

    @Override
    public Object[] toArray() {

        Object[] array = new Object[size];
        MyTreeSetIterator iter = new MyTreeSetIterator();

        for (int i = 0; i < size; i++) {
            array[i] = iter.next();
        }

        return array;
    }

    // TODO: 16.09.2016 toArray[]
   @Override
    public <T> T[] toArray(T[] a) {
        return null;
    }

    @Override
    public boolean add(E value) {
        if (root == null) {
            root = new TreeNode(value, null, null, null);
            size++;
            return true;
        }

        TreeNode curr = root;

        while (curr != null) {
            int compare = canCompare(value, (E) curr.value);
            if (compare < 0) {

                if (curr.left == null) {
                    curr.left = new TreeNode(value, curr, null, null);
                    size++;
                    return true;
                }

                curr = curr.left;
            } else if (compare > 0) {

                if (curr.right == null) {
                    curr.right = new TreeNode(value, curr, null, null);
                    size++;
                    return true;
                }

                curr = curr.right;
            } else {
                return false;
            }
        }
        return false;
    }

    @Override
    public boolean remove(Object o) {

        TreeNode forRemove = findNode(root, (E) o);
        if (forRemove == null){
            return false;
        }

        if (forRemove.left == null && forRemove.right == null) {
            if (isChildLeft(forRemove)){
                forRemove.parent.left = null;
            } else {
                forRemove.parent.right = null;
            }
            size--;
            return true;
        }

        if ((forRemove.left != null && forRemove.right == null)){
            if (isChildLeft(forRemove)){
                forRemove.parent.left = forRemove.left;
            } else {
                forRemove.parent.right = forRemove.left;
            }
            forRemove.left.parent = forRemove.parent;
            size--;
            return true;
        }

        if ((forRemove.left == null && forRemove.right != null)){
            if (isChildLeft(forRemove)){
                forRemove.parent.left = forRemove.right;
            } else {
                forRemove.parent.right = forRemove.right;
            }
            forRemove.right.parent = forRemove.parent;
            size--;
            return true;
        }

        if (forRemove.left != null && forRemove.right != null){

            if (isChildLeft(forRemove)){
                forRemove.parent.left = forRemove.right;
            } else {
                forRemove.parent.right = forRemove.right;
            }
            forRemove.right.parent = forRemove.parent;

            TreeNode newParentForLeftBranch = getLeftTailNode(forRemove.right);
            newParentForLeftBranch.left = forRemove.left;
            forRemove.left.parent = newParentForLeftBranch;
            size--;
            return true;
        }
        return false;
    }

    private boolean isChildLeft(TreeNode curr) {
        return curr.parent.left == curr;
    }

    // TODO: 16.09.2016 containsAll
    @Override
    public boolean containsAll(Collection<?> c) {
        return false;
    }

    // TODO: 16.09.2016 AddAll
    @Override
    public boolean addAll(Collection<? extends E> c) {
        return false;
    }

    // TODO: 16.09.2016 retainAll
    @Override
    public boolean retainAll(Collection<?> c) {
        return false;
    }

    // TODO: 16.09.2016 removeAll
    @Override
    public boolean removeAll(Collection<?> c) {
        return false;
    }

   @Override
    public void clear() {
        root = null;
        size = 0;
    }

    private static class TreeNode<E> {

        E value;
        TreeNode parent;
        TreeNode left;
        TreeNode right;

        public TreeNode(E value, TreeNode parent,
                        TreeNode left, TreeNode right) {
            this.value = value;
            this.parent = parent;
            this.left = left;
            this.right = right;
        }
    }

    private class MyTreeSetIterator<E> implements Iterator<E> {
        private TreeNode curr;

        private MyTreeSetIterator() {
            curr = getLeftTailNode(root);
        }

        @Override
        public boolean hasNext() {
            return curr != null;
        }

        @Override
        public E next() {

            TreeNode current = curr;

            if (curr.right != null) {
                curr = getLeftTailNode(curr.right);
                return (E) current.value;

            } else while (true) {
                if (curr.parent == null) {
                    curr = null;
                    return (E) current.value;
                }
                if (curr.parent.left == curr) {
                    curr = curr.parent;
                    return (E) current.value;
                }
                curr = curr.parent;
            }
        }
    }

    private TreeNode getLeftTailNode(TreeNode root) {

        TreeNode curr = root;

        while (curr.left != null) {
            curr = curr.left;
        }

        return curr;
    }

    @Override
    public String toString() {
        return "MyTreeSetInteger{" +
                "root=" + root.value +
                ", size=" + size +
                '}' + "\n" +
                treeToString(root);
    }

    private String treeToString(TreeNode root) {
        StringBuilder result = new StringBuilder();

        if(root == null) {
            return "-";
        }
        return result.append("(")
                .append(treeToString(root.left))
                .append(root.value)
                .append(treeToString(root.right))
                .append(")").toString();
    }
}