数据结构之TreeMap

package datastruct.binarytree;

import java.util.Comparator;

public class HtqTreeMap<K,V> {
    private static class Node<K,V>{
        boolean color;
        K key;
        V value;
        Node<K,V> parent;
        Node<K,V> left;
        Node<K,V> right;

        private Node(K key, V value, Node<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        private boolean hasTwoChildren(){
            return left!=null&&right!=null;
        }

        private boolean isLeaf(){
            return left==null&&right==null;
        }

        private boolean isLeftChild(){
            return (parent==null)?false:parent.left==this;
        }

        private boolean isRightChild(){
            return (parent==null)?false:parent.right==this;
        }

        private Node<K,V> getUncleNode(){
            if (parent==null) return null;
            if (parent.isLeftChild()){
                return parent.parent.right;
            }else if (parent.isRightChild()){
                return parent.parent.left;
            }else {
                return null;
            }
        }

        private Node<K,V> getBrotherNode(){
            if (this.isLeftChild()){
                return parent.right;
            }else if (this.isRightChild()){
                return parent.left;
            }else {
                return null;
            }
        }

        private Node<K,V> blackColor(){
            this.color = true;
            return this;
        }

        private Node<K,V> redColor(){
            this.color = false;
            return this;
        }
    }

    private int size;

    private Node<K,V> root;

    private Comparator<K> comparator;

    public HtqTreeMap() {
        this(null);
    }

    public HtqTreeMap(Comparator<K> comparator) {
        this.comparator = comparator;
    }

    public int size(){
        return size;
    }

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

    public boolean contains(K key){
        return node(key)!=null;
    }

    private Node<K,V> node(K key) {
        if (key==null) return null;

        Node<K,V> node = root;
        while (node!=null){
            int result = compare(key,node.key);
            if (result<0){
                node = node.left;
            }else if (result>0){
                node = node.right;
            }else {
                return node;
            }
        }
        return null;
    }

    private int compare(K k1, K k2) {
        return (comparator!=null)?comparator.compare(k1,k2):((Comparable<K>)k1).compareTo(k2);
    }

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

    public void put(K key,V value){
        if (key==null) return;

        if (root==null){
            root = new Node<>(key,value,null);
            size++;
            afterPut(root);
            return;
        }

        Node<K,V> node = root;
        Node<K,V> parent = null;
        int result = 0;
        while (node!=null){
            parent = node;
            result = compare(key,node.key);
            if (result<0){
                node = node.left;
            }else if (result>0){
                node = node.right;
            }else {
                node.key = key;
                node.value = value;
                return;
            }
        }

        Node<K,V> newNode = new Node<>(key,value,parent);
        if (result<0){
            parent.left = newNode;
        }else {
            parent.right = newNode;
        }
        size++;
        afterPut(newNode);
    }

    private void afterPut(Node<K, V> node) {
        if (node.parent==null){
            node.blackColor();
            return;
        }

        if (isBlack(node.parent)){
            return;
        }

        Node<K, V> uncle = node.getUncleNode();
        if (isRed(uncle)){
            node.parent.blackColor();
            uncle.blackColor();
            afterPut(node.parent.parent.redColor());
        }else {
            Node<K, V> grand = node.parent.parent;
            Node<K, V> parent = node.parent;

            if (parent.isLeftChild()){
                if (node.isLeftChild()){
                    grand.redColor();
                    parent.blackColor();
                    rightRotate(grand);
                }else {
                    node.blackColor();
                    grand.redColor();
                    leftRotate(parent);
                    rightRotate(grand);
                }
            }else {
                if (node.isLeftChild()){
                    node.blackColor();
                    grand.redColor();
                    rightRotate(parent);
                    leftRotate(grand);
                }else {
                    grand.redColor();
                    parent.blackColor();
                    leftRotate(grand);
                }
            }
        }

    }

    private void rightRotate(Node<K, V> grand) {
        Node<K, V> parent = grand.left;
        Node<K, V> node = parent.right;

        grand.left = node;
        parent.right = grand;

        parent.parent = grand.parent;
        if (grand.isLeftChild()){
            grand.parent.left = parent;
        }else if (grand.isRightChild()){
            grand.parent.right = parent;
        }else {
            root = parent;
        }

        if (node!=null){
            node.parent = grand;
        }

        grand.parent = parent;
    }

    private void leftRotate(Node<K, V> grand) {
        Node<K, V> parent = grand.right;
        Node<K, V> node = parent.left;

        grand.right = node;
        parent.left = grand;

        parent.parent = grand.parent;
        if (grand.isLeftChild()){
            grand.parent.left = parent;
        }else if (grand.isRightChild()){
            grand.parent.right = parent;
        }else {
            root = parent;
        }

        if (node!=null){
            node.parent = grand;
        }

        grand.parent = parent;
    }

    //error
    private boolean isBlack(Node<K,V> node){
        return node==null||node.color;
    }

    //error
    private boolean isRed(Node<K,V> node){
        return node!=null&&!node.color;
    }

    public V remove(K key){
        Node<K,V> node = node(key);
        if (node==null) return null;
        V v = node.value;

        if (node.hasTwoChildren()){
            Node<K,V> post = postNode(node);
            node.key = post.key;
            node.value = post.value;
            node = post;
        }

        if (node.isLeaf()){
            if (node.isLeftChild()){
                node.parent.left = null;
            }else if (node.isRightChild()){
                node.parent.right = null;
            }else {
                root = null;
            }
        }else {
            Node<K,V> replace = (node.left==null)?node.right:node.left;

            replace.parent = node.parent;
            if (node.isLeftChild()){
                node.parent.left = replace;
            }else if (node.isRightChild()){
                node.parent.right = replace;
            }else {
                root = replace;
            }
            node = replace;
        }
        size--;
        afterRemove(node);
        return v;
    }

    private void afterRemove(Node<K, V> node) {
        if (isRed(node)){
            node.blackColor();
            return;
        }

        //error
        if (node.parent==null){
            node.blackColor();
            return;
        }

        //error
        Node<K,V> brother = (node.parent.left==null||node.parent.left==node)?node.parent.right:node.parent.left;
        if (brother.isLeftChild()){
            if (isRed(brother)){
                brother.blackColor();
                node.parent.redColor();
                rightRotate(node.parent);
                brother = node.parent.left;
            }

            boolean pColor = node.parent.color;

            if (isRed(brother.left)||isRed(brother.right)){
                if (isBlack(brother.left)&&isRed(brother.right)){
                    leftRotate(brother);
                    brother = node.parent.left;
                }

                brother.color = pColor;
                node.parent.blackColor();
                brother.left.blackColor();
                rightRotate(node.parent);

            }else {//error
                node.parent.blackColor();
                brother.redColor();
                if (pColor){
                    afterRemove(node.parent);
                }
            }
        }else {
            if (isRed(brother)){
                brother.blackColor();
                node.parent.redColor();
                leftRotate(node.parent);
                brother = node.parent.right;
            }

            boolean pColor = node.parent.color;

            if (isRed(brother.left)||isRed(brother.right)){
                if (isBlack(brother.right)&&isRed(brother.left)){
                    rightRotate(brother);
                    brother = node.parent.right;
                }

                brother.color = pColor;
                node.parent.blackColor();
                brother.right.blackColor();
                leftRotate(node.parent);

            }else {//error
                node.parent.blackColor();
                brother.redColor();
                if (pColor){
                    afterRemove(node.parent);
                }
            }
        }
    }

    private Node<K,V> preNode(Node<K,V> node){
        Node<K,V> left = node.left;
        if (left!=null){
            while (left.right!=null){
                left = left.right;
            }
            return left;
        }

        Node<K,V> parent = node.parent;
        while (parent!=null&&parent.isLeftChild()){
            parent = parent.parent;
        }
        return parent;
    }

    private Node<K,V> postNode(Node<K,V> node){
        Node<K,V> right = node.right;
        if (right!=null){
            while (right.left!=null){
                right = right.left;
            }
            return right;
        }

        Node<K,V> parent = node.parent;
        while (parent!=null&&parent.isRightChild()){
            parent = parent.parent;
        }
        return parent;
    }

    public String toString(){
        StringBuilder stringBuilder = new StringBuilder();
        toString(root,stringBuilder,"");
        return stringBuilder.toString();
    }

    private StringBuilder toString(Node<K,V> node,StringBuilder stringBuilder,String fix){
        if (node==null) return stringBuilder;
        toString(node.left,stringBuilder,fix+"L-->");
        stringBuilder.append(fix);
        if (node.color){
            stringBuilder.append("["+node.key+"_"+node.value+"]").append("\n");
        }else {
            stringBuilder.append("[R_"+node.key+"_"+node.value+"]").append("\n");
        }
        toString(node.right,stringBuilder,fix+"R-->");
        return stringBuilder;
    }
}