数据结构之二叉平衡搜索树（RBT）

package datastruct.binarytree;

import java.util.Comparator;
import java.util.Objects;

public class RBTree<E> {
    private static class Node<E>{
        boolean color;
        E item;
        Node<E> parent;
        Node<E> left;
        Node<E> right;

        private Node(E item, Node<E> parent) {
            this.item = item;
            this.parent = parent;
        }

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

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

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

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

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

        //error uncle left/right !!
        private Node<E> 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 static final boolean BLACK = true;

    private static final boolean RED = false;

    private int size;

    private Node<E> root;

    private Comparator<E> comparator;

    public RBTree() {
        this(null);
    }

    public RBTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public int size(){
        return size;
    }

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

    public boolean contains(E e){
        return node(e)!=null;
    }

    private Node<E> node(E e) {
        if (e==null) return null;

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

        return null;
    }

    private int compare(E e1, E e2) {
        return (comparator!=null)?comparator.compare(e1,e2):((Comparable<E>)e1).compareTo(e2);
    }

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

    public void add(E e){
        if (e==null) return;

        if (root==null){
            root = new Node<>(e,null);
            size++;
            afterAdd(root);
            return;
        }

        Node<E> node = root;
        Node<E> parent = null;
        int result = 0;
        while (node!=null){
            parent = node;
            result = compare(e,node.item);
            if (result<0){
                node = node.left;
            }else if (result>0){
                node = node.right;
            }else {
                node.item = e;
                return;
            }
        }

        Node<E> newNode = new Node<>(e,parent);
        if (result<0){
            parent.left = newNode;
        }else if (result>0){
            parent.right = newNode;
        }

        size++;

        afterAdd(newNode);
    }

    private void afterAdd(Node<E> node) {
        if (node.parent==null){
            blackColor(node);
            return;
        }

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

        Node<E> uncle = node.getUncleNode();
        if (isRed(uncle)){//error
            blackColor(uncle);
            blackColor(node.parent);
            afterAdd(redColor(node.parent.parent));
        }else {
            Node<E> parent = node.parent;
            Node<E> grand = parent.parent;
            if (parent.isLeftChild()){
                if (node.isLeftChild()){
                    blackColor(parent);
                    redColor(grand);
                    rightRotate(grand);
                }else {
                    blackColor(node);
                    redColor(grand);
                    leftRotate(parent);
                    rightRotate(grand);
                }
            }else {
                if (node.isLeftChild()){
                    blackColor(node);
                    redColor(parent);
                    redColor(grand);
                    rightRotate(parent);
                    leftRotate(grand);
                }else {
                    blackColor(parent);
                    redColor(grand);
                    leftRotate(grand);
                }
            }
        }
    }

    private void rightRotate(Node<E> grand) {
        Node<E> parent = grand.left;
        Node<E> right = parent.right;

        grand.left = right;
        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 (right!=null){
            right.parent = grand;
        }

        grand.parent = parent;

    }

    //error grand. left/right
    private void leftRotate(Node<E> grand) {
        Node<E> parent = grand.right;
        Node<E> left = parent.left;

        grand.right = left;
        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 (left!=null){
            left.parent = grand;
        }

        grand.parent = parent;
    }

    private Node<E> blackColor(Node<E> node){
        node.color = true;
        return node;
    }

    private Node<E> redColor(Node<E> node){
        if (node==null) throw new IllegalArgumentException();
        node.color = false;
        return node;
    }

    private boolean isBlack(Node<E> node){
        return (node==null)||node.color;
    }

    private boolean isRed(Node<E> node){
        return (node!=null)&&!node.color;
    }

    public void remove(E e){
        Node<E> node = node(e);

        if (node==null) return;

        if (node.hasTwoChildren()){
            Node<E> post = postNode(node);
            node.item = post.item;
            node = post;
        }

        if (node.isLeaf()){
            if (node.isLeftChild()){
                node.parent.left = null;
            }else if (node.isRightChild()){
                node.parent.right = null;
            }else {
                root = null;
            }
        }else {//error no deal parent replace.parent = node.parent;
            Node<E> 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--;
        afterDelete(node);
    }

    private void afterDelete(Node<E> node) {
        if (isRed(node)){
            blackColor(node);
            return;
        }

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

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

                brother.color = pColor;
                isBlack(node.parent);
                isBlack(brother.left);
                rightRotate(node.parent);
            }else {
                boolean pColor = node.parent.color;
                blackColor(node.parent);
                redColor(brother);
                if (pColor){
                    afterDelete(node.parent);
                }
            }
        }else {
            if (isRed(brother)){
                blackColor(brother);
                redColor(node.parent);
                leftRotate(node.parent);
                brother = node.parent.right;
            }

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

                brother.color = pColor;
                isBlack(node.parent);
                isBlack(brother.right);
                leftRotate(node.parent);
            }else {
                boolean pColor = node.parent.color;
                blackColor(node.parent);
                redColor(brother);
                if (pColor){
                    afterDelete(node.parent);
                }
            }
        }
    }

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

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

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

        Node<E> 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<E> 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.item+"]").append("\n");
        }else {
            stringBuilder.append("[[R]_"+node.item+"]").append("\n");
        }
        toString(node.right,stringBuilder,fix+"R-->");
        return stringBuilder;
    }
}