rbt

package rbtmap;

public class RedBlackTreeMap<K extends Comparable<K>, V> implements MapInterface<K, V> {

    private Node<K, V> root = null;

    @Override
    public V getValue(K key) {
        Node<K, V> tmp = findNode(key, root);
        if (tmp == null)
            return null;
        else
            return tmp.getVal();
    }

    @Override
    public void setValue(K key, V value) {
        if (key == null)
            throw new IllegalArgumentException("Given key is null!");
        if (root == null) {
            root = new Node<> (key, value);
            root.paintBlack();
        }
        else {
            Node<K, V> tmp = findNode(key, root);
            if (tmp != null)
                tmp.setVal(value);
            else {
                addNode(key, value, root);
            }
        }
    }

    private Node<K, V> findNode(K key, Node<K, V> root) {
        if (key.compareTo(root.getKey()) > 0)
            findNode(key, root.getRightChild());
        if (key.compareTo(root.getKey()) < 0)
            findNode(key, root.getLeftChild());
        return key == root.getKey()? root : null;
    }

    private Node<K,V> addNode(K key, V value, Node<K,V> subRoot) {
        if (subRoot != null) {
            if (subRoot.getKey().compareTo( key ) > 0)
                subRoot = addNode(key, value, subRoot.getLeftChild());
            else
                subRoot = addNode(key, value, subRoot.getRightChild());
            fixTreeAfterInsertion(subRoot);
        }
        return new Node<>(key, value);
    }

    private void fixTreeAfterInsertion(Node<K, V> pivot) {
        if (pivot.getLeftChild().isRed() && pivot.getRightChild().isRed()) {
            pivot.getRightChild().paintBlack();
            pivot.getLeftChild().paintBlack();
            pivot.paintRed();
        }
        if (pivot.getRightChild().isRed() && !pivot.getLeftChild().isRed()) {
            pivot = rotateLeft(pivot);
            pivot.getLeftChild().paintRed();
            pivot.paintBlack();
        }
        if (pivot.getLeftChild().isRed() && pivot.getLeftChild().getLeftChild().isRed()) {
            pivot.getLeftChild().paintBlack();
            pivot = rotateRight(pivot);
            pivot.paintRed();
        }
    }

    private Node<K,V> rotateLeft(Node<K,V> pivot) {
        Node<K,V> left = pivot.getLeftChild();
        Node<K,V> right = pivot.getRightChild();
        right.setLeftChild(pivot);
        pivot.setRightChild(left);
        return right;
    }

    private Node<K,V> rotateRight(Node<K,V> pivot) {
        Node<K,V> left = pivot.getLeftChild();
        Node<K,V> right = pivot.getRightChild();
        left.setRightChild(pivot);
        pivot.setLeftChild(right);
        return left;
    }
}

/*
Node.java
*/

package rbtmap;

public class Node<K extends Comparable<K>, V> {
    private K key;
    private V value;
    private Node<K, V> leftChild;
    private Node<K, V> rightChild;
    private enum Color {
        RED, BLACK
    }
    private Color color;

    private Color getColor() {
        return color;
    }

    private void setColor(Color color) {
        this.color = color;
    }

    public K getKey() {
        return key;
    }

    public V getVal() {
        return value;
    }

    public Node<K, V> getLeftChild() {
        return leftChild;
    }

    public Node<K, V> getRightChild() {
        return rightChild;
    }

    public void setLeftChild(Node<K, V> leftChild) {
        this.leftChild = leftChild;
    }

    public void setRightChild(Node<K, V> rightChild) {
        this.rightChild = rightChild;
    }

    public void setVal(V value) {
        this.value = value;
    }

    public boolean isRed() {
        return getColor() == Color.RED;
    }

    public void paintBlack() {
        setColor(Color.BLACK);
    }

    public void paintRed() {
        setColor(Color.RED);
    }

    public Node(K key, V value) {
        this.key = key;
        this.value = value;
        leftChild = null;
        rightChild = null;
        setColor(Color.RED);
    }
}