用红黑树实现Map

直接看代码：

/**
 * 直接把TreeMap当作一颗红黑树去实现，实现一个Map
 * @param <K>
 * @param <V>
 */
@SuppressWarnings({"all"})
public class TreeMap<K,V> implements Map<K,V>{

    private static final boolean RED = false ;

    private static final boolean BLACK = true ;

    private int size ;
    //确定根节点
    private Node<K,V> root ;

    //传入一个Comparator对象
    private Comparator<K> comparator ;

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

    public TreeMap() {
        //调用的是本类的有参构造器，只不过传入的是空参
        this(null) ;
    }
    /**
     * 树节点元素的数量
     */
    public int size() {
        return size;
    }
    /**
     * 是否为空
     */
    public boolean isEmpty() {
        return size == 0 ;
    }
    /**
     * 清空所有元素
     */
    public void clear() {
        root = null ;
        size = 0 ;
    }
    /**
     * @return 返回值返回的是被覆盖掉的value值是什么
     */
    @Override
    public V put(K key, V value) {
        //key不可以为空
        elementNotNullCheck(key);

        //1.添加的是第一个节点
        if (root == null) {
            //对于根节点而言没有父节点
            root = new Node<>(key,value,null ) ;
            size ++ ;
            //添加root节点之后对其进行调整
            afterPut(root);
            return null;
        }
        //2.说明添加的不是树中第一个元素
        //2.1 确定根节点 root
        Node<K,V> node = root ;
        //2.2 假设出待插入节点对应的父节点的值
        Node<K,V> parent = null ;
        //2.3 记录最后一次比较器返回的值 来确定待插入节点插入到哪个方向
        int cmp = 0 ;
        //2.4 循环寻找待插入节点需要插入到的位置
        while (node != null) {
            //记录cmp
            cmp = compare(node.key,key) ;
            //记录parent
            parent = node ;
            if (cmp > 0) {
                //说明已存在节点元素值大，那么插到左边
                node = node.left ;
            } else if (cmp < 0) {
                node = node.right ;
            } else {
                //cmp == 0 如果自定义比较器比较出相等时，我们这里的处理方式是覆盖值操作
                node.key = key ;
                V oldValue = node.value ;
                node.value = value ;
                return oldValue ;//返回的是被覆盖的value值
            }
        }
        //3.退出循环之后，根据记录的parent值 进行创建待插入节点对应的Node对象
        Node<K,V> newNode = new Node<>(key,value,parent) ;
        //4.根据记录的cmp值和parent值，进行确定把待插入节点插入到哪一个位置
        if (cmp > 0) {
            parent.left = newNode ;
        } else {
            //cmp < 0 的情况，cmp == 0 直接就返回了已经
            parent.right = newNode ;
        }
        size ++ ;
        //添加node节点后 进行的调整
        afterPut(node); ;
        return null ;
    }

    /**
     * 添加操作可能毁坏红黑树的性质
     * afterPut修复红黑树的性质
     * @param node
     */
    private void afterPut(Node<K, V> node) {
        Node<K,V> parent = node.parent ;
        //1.情况一中的四种添加位置，是无需进行修复处理红黑树的。
        //1.1添加的是根节点
        if (parent == null) {
            black(node) ;//根节点直接染成黑色
            return;
        }
        //1.2如果父节点是黑色，直接返回
        if (isBlack(parent)) return ;
        //2.情况二中的八种添加位置，是需要进行修复处理红黑树的。
        //2.1拿到uncle节点进行判断 uncle节点：叔父节点即是父节点的兄弟节点
        Node<K,V> uncle = parent.sibling() ;
        //2.2祖父节点
        Node<K,V> grand = parent.parent ;
        //2.3 叔父节点为RED，只需进行染色 然后进行递归grand作为新添加的节点进行处理
        if (isRed(uncle)) {
            black(parent) ;
            black(uncle) ;
            //把祖父节点当作是新添加的节点进行处理
            red(grand) ;
            afterPut(grand); ;
            return;
        }
        //2.4 叔父节点不为RED,需要进行旋转处理
        if (parent.isLeftChild()) { //parent是grand的左子树节点 L
            if (node.isLeftChild()) { //node是parent的左子树节点 LL
                //LL和RR为一类，parent染成BLack，grand染成RED
                black(parent) ;
                red(grand) ;
                rotateRight(grand);
            } else { //LR
                //LR 和 RL为一类，新添加的节点染成Black，grand染成RED
                black(node) ;
                red(grand) ;
                rotateLeft(parent) ;
                rotateRight(grand) ;
            }
        } else { //R
            if (node.isLeftChild()) { //RL
                //LR 和 RL为一类，新添加的节点染成Black，grand染成RED
                black(node) ;
                red(grand) ;
                rotateRight(parent); ;
                rotateLeft(grand); ;
            } else {//RR
                //LL和RR为一类，parent染成BLack，grand染成RED
                black(parent) ;
                red(grand) ;
                rotateLeft(grand);
            }
        }
    }
    private int compare(K k1,K k2)  {
        if (comparator != null) {
            return comparator.compare(k1,k2) ;
        }
        //如果比较器为空，那么要进行强转化为Comparable类型进行使用对应比较逻辑
        return ((Comparable<K>)k1).compareTo(k2) ;
    }
    private void elementNotNullCheck(K key) {
        if (key == null) {
            throw new IllegalArgumentException("element must not be null") ;
        }
    }
    @Override
    public V get(K key) {
        Node<K,V> node = node(key) ;
        return node != null ? node.value : null;
    }

    @Override
    public V remove(K key) {
        return remove(node(key));
    }

    @Override
    public boolean containsKey(K key) {
        return node(key) != null ;
    }

    /**
     * 由于只有Key具有可比较性，所以寻找value看是否存在时
     * 我们需要遍历整个红黑树
     * 遍历时使用的是层序遍历
     */
    @Override
    public boolean containsValue(V value) {
        Queue<Node<K,V>> queue = new LinkedList<>() ;
        queue.offer(root) ;
        while (!queue.isEmpty()) {
            //队头出队
            Node<K,V> node = queue.poll() ;
            if (valEquals(node.value,value)) {
                return true ;//值相等 说明找到了 返回true
            }
            if (node.left != null) {
                queue.offer(node.left) ;
            }
            if (node.right != null) {
                queue.offer(node.right) ;
            }
        }
        return false ;
    }

    /**
     * 遍历 使用中序遍历：得出从小到大的结果 !
     * @param visitor
     */
    @Override
    public void traversal(Visitor<K, V> visitor) {
        if (visitor == null) {
            return;
        }
        traversal(root,visitor) ;
    }

    private void traversal(Node<K,V> node,Visitor<K,V> visitor) {
        if (node == null || visitor.stop) {
            return;
        }
        traversal(node.left,visitor) ;
        if (visitor.stop) {
            return;
        }
        visitor.visit(node.key,node.value) ;
        traversal(node.right,visitor) ;
    }
    /**
     * 比较v1与v2是否相等
     */
    private boolean valEquals(V v1,V v2) {
        return v1 == null ? v2 == null : v1.equals(v2) ;
    }
    private V remove(Node<K,V> node) {
        if (node == null) {
            return null;
        }
        V oldValue = node.value ;
        size -- ;
        //(1)删除的是度为2的节点
        if (node.hasTwoChildren()) {
            //找到该节点对应的前驱或后继节点
            Node<K,V> s = succeed(node) ;//找到后继节点
            //使用后继节点的值进行覆盖度为2的节点的值
            node.key = s.key ;
            node.value = s.value ;
            //把待删除的节点进行赋值为前驱或后继节点，使用后面的逻辑进行删除前驱后继节点
            node = s ;
        }
        //必须搞出一个replacement进行记录删除的是根节点左子树还是右子树的节点
        Node<K,V> replacement = node.left != null ? node.left : node.right ;
        //(2)删除的是度为1的节点
        if (replacement != null) {
            //统一进行更新删除节点度为1的操作之后的parent指向值
            replacement.parent = node.parent ;
            //2.1 如果是度为1的节点并且是root根节点
            if (node.parent == null) {
                root = replacement ;
            } else if (node == node.parent.left) {
                //2.2 如果是左子树上的度为1的节点
                node.parent.left = replacement ;
            } else if (node == node.parent.right) {
                //2.3 如果是右子树上的度为1的节点
                node.parent.right = replacement ;
            }
            //真正被删除的是度为1或2的前驱或后继节点
            afterRemove(node,replacement) ;
        } else if (node.parent == null) {// (3) 如果删除的是叶子节点并且为root根节点
            root = null ;
        } else {//(4) 如果删除的是叶子节点但是不为root根节点
            if (node == node.parent.right) {
                //4.1 如果是左子树上的度为0的节点
                node.parent.right = null ;
            } else {//node = node.parent.left
                //4.2 如果是右子树上的度为0的节点
                node.parent.left = null ;
            }
        }
        return oldValue ;
    }
    private void afterRemove(Node<K,V> node, Node<K,V> replacement) {
        //1。如果删除的节点是红色RED，删除之后无需处理
        if (isRed(node)) return;
        //2.用以取代被删除的node的子节点是RED时
        if (isRed(replacement)) {
            black(replacement) ;
            return;
        }
        Node<K,V> parent = node.parent ;
        //删除的是根节点
        if (parent == null) return;
        //3.删除的是黑色叶子节点
        //判断被删除的node是左还有右:
        //parent.left == null说明被删除的node节点是parent的左子树，
        // 那么被删除的node节点的兄弟节点即是parent的右子树，反之同理即可
        boolean left = parent.left == null ;
        Node<K,V> sibling = left ? parent.right : parent.left ;
        if (left) {//3.1被删除的节点在左边，那么兄弟节点sibling在右边
            if (isRed(sibling)) {
                //3.2.1删除的是黑色叶子节点并且sibling是红色，那么转换为sibling为黑色
                black(sibling) ;
                red(parent) ;
                rotateRight(parent) ;
                //更换兄弟
                sibling = parent.right ;
            }
            //3.2.2删除的是黑色叶子节点并且sibling为BLACK
            //3.2.2.1 兄弟节点没有一个是红色子节点，父节点要下根兄弟节点合并
            if (isBlack(sibling.right) && isBlack(sibling.left)) {
                //记录一下是否父节点也会下溢
                boolean parentBlack = isBlack(parent) ;
                black(parent) ;
                red(sibling) ;
                if (parentBlack) {
                    afterRemove(parent,null) ;
                }
            } else {  //3.2.2.2 兄弟节点至少有1个红色子节点，被删除节点向兄弟节点借用元素
                //兄弟节点的左边是黑色,说明右边是红色子节点
                if (isBlack(sibling.right)) {
                    rotateLeft(sibling) ;
                    sibling = parent.right ;//旋转后要更新兄弟节点
                    color(sibling,colorOf(parent)) ;
                    black(sibling.right) ;
                    black(parent) ;
                }
                //兄弟节点的左边是红色,说明左边是红色子节点
                color(sibling,colorOf(parent)) ;
                black(sibling.right) ;
                black(parent) ;
                rotateRight(parent);
            }
        } else {//3.2被删除的节点在右边，那么兄弟节点sibling在左边
            if (isRed(sibling)) {
                //3.2.1删除的是黑色叶子节点并且sibling是红色，那么转换为sibling为黑色
                black(sibling) ;
                red(parent) ;
                rotateRight(parent) ;
                //更换兄弟
                sibling = parent.left ;
            }
            //3.2.2删除的是黑色叶子节点并且sibling为BLACK
            //3.2.2.1 兄弟节点没有一个是红色子节点，父节点要下根兄弟节点合并
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                //记录一下是否父节点也会下溢
                boolean parentBlack = isBlack(parent) ;
                black(parent) ;
                red(sibling) ;
                if (parentBlack) {
                    afterRemove(parent,null) ;
                }
            } else {  //3.2.2.2 兄弟节点至少有1个红色子节点，被删除节点向兄弟节点借用元素
                //兄弟节点的左边是黑色,说明右边是红色子节点
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling) ;
                    sibling = parent.left ;//旋转后要更新兄弟节点
                    color(sibling,colorOf(parent)) ;
                    black(sibling.left) ;
                    black(parent) ;
                }
                //兄弟节点的左边是红色,说明左边是红色子节点
                color(sibling,colorOf(parent)) ;
                black(sibling.left) ;
                black(parent) ;
                rotateRight(parent);
            }
        }
    }
    /**
     * 找到node节点对应的前驱节点：
     * 即是二叉树上的元素进行中序遍历顺序之后得出的一串数字节点中，node节点所在其中位置的前一个节点
     * @param node
     * @return
     */
    private Node<K,V> predecessor(Node<K,V> node) {
        if (node == null) {
            return null ;
        }
        Node<K,V> p = node ;
        if (p.left != null) {
            p = p.left ;
            while (p.right != null) {
                p = p.right ;
            }
            //说明p==null，说明一个右子节点为空，
            // 那么最后一次不为空的node.right即为node的前驱节点
            return p ;
        }
        //说明p.left == null，没有进入if语句
        //我们就要找父节点
        //条件为：父节点不为空，并且。。
        while ( p.parent != null &&  p == p.parent.left) {
            p = p.parent ;
        }
        //到这里 有两种情况：
        //1.p.parent == null (说明该node没有前驱节点 返回p.parent(null)即可)
        //2. p == p.parent.right(找到该node的前驱节点，返回前驱节点：p.parent)
        return p.parent ;
    }

    /**
     * 找到后继节点
     * @param node
     * @return
     */
    private Node<K,V> succeed(Node<K,V> node) {
        if (node == null) {
            return null ;
        }
        //1,
        Node<K,V> p = node ;
        if (p.right != null) {
            p = p.right ;
            while (p.left != null) {
                p = p.left ;
            }
            return p ;
        }
        //p.right == null
        //目的是找node节点的后继节点，所以要找第一个比它大的
        //所以在父节点不为空 并且一直在变化的父节点的右子树 、
        // 如果它发现在左子树后，那么就要退出循环，说明找到了第一个比它大的节点
        if (p.parent != null && p == p.parent.right) {
            p = p.parent ;
        }
        //p.parent == null
        //p == p.parent.left
        return p.parent ;
    }
    /**
     * 对传入的节点进行染色处理
     * @param node 想要染色的节点
     * @param color 把该节点node染成color颜色
     * @return 返回染色后的节点
     */
    private Node<K,V> color(Node<K,V> node, boolean color) {
        if (node == null) {
            return node ;
        }
        node.color = color ;
        return node ;
    }

    /**
     * 把该node节点染成RED
     */
    private Node<K,V> red(Node<K,V> node) {
        return color(node,RED) ;
    }

    /**
     * 把该node节点染成BLACK
     */
    private Node<K,V> black(Node<K,V> node) {
        return color(node,BLACK) ;
    }

    /**
     * 传入一个node节点，进行返回该节点是什么颜色
     */
    private boolean colorOf(Node<K,V> node) {
        //如果node为null，即是null节点 即是BLACK
        // 【对null节点进行特殊处理的原因是：我们一开始并没有对null节点的颜色进行设置】
        //如果不为空，那么返回node实际的颜色即可
        return node == null ? BLACK : node.color ;
    }

    /**
     * 判断该node节点颜色color是否为BLACK
     */
    private boolean isBlack(Node<K,V> node) {
        return colorOf(node) == BLACK ;
    }

    /**
     * 判断该node节点颜色color是否为RED
     */
    private boolean isRed(Node<K,V> node) {
        return colorOf(node) == RED ;
    }

    /**
     * 左旋转
     * @param grand 想要进行旋转的节点,传过来的是什么，什么节点就进行左旋转
     */
    private void rotateLeft(Node<K,V> grand) {
        //进行左旋转，说明是RR，可以确定parent
        Node<K,V> parent = grand.right ;
        //child即是T1
        Node<K,V> child = parent.left ;
        grand.right =  child; //1.
        parent.left = grand ; //2.
        //执行旋转之后的逻辑
        afterRotate(grand,parent,child);
    }
    /**
     * 右旋转
     * @param grand 想要进行旋转的节点，传过来的是什么，什么节点就进行右旋转
     */
    private void rotateRight(Node<K,V> grand) {
        //1.进行右旋转，说明是LL，那么可以确定parent
        Node<K,V> parent = grand.left ;
        //2.找出T2
        Node<K,V> child = parent.right ;
        //3.对g执行右旋转
        grand.left = child ; //3.1
        parent.right = grand ;//3.2
        //执行旋转之后的逻辑
        afterRotate(grand,parent,child);
    }
    /**
     * 旋转之后统一要做的逻辑
     * @param grand （最低的失去平衡的节点）
     * @param parent (parent是grand的子节点)
     * @param child 旋转过程中父节点改变的节点（T1或T2）
     */
    private void afterRotate(Node<K,V> grand, Node<K,V> parent, Node<K,V> child) {
        //3.让parent成为这颗子树的根节点
        //3.1parent先连上grand的父节点
        parent.parent = grand.parent ;
        //3.2 grand再去连上parent
        if (grand.isLeftChild()) {
            //如果grand是父节点的左子树
            grand.parent.left = parent ;
        } else if (grand.isRightChild()) {
            //如果grand是父节点的右子树
            grand.parent.right = parent ;
        } else {//既不是父节点的左子树又不是右子树，那么说明grand即是root根节点
            root = parent ;
        }
        //4.更新child(T1)的parent
        if (child != null) {
            child.parent = grand ;
        }
        //5.更新grand的parent
        grand.parent = parent ;
    }

    /**
     * 找到相同的key 返回对应的Node节点
     * @param key
     * @return
     */
    private Node<K,V> node(K key) {
        Node<K,V> node = root ;
        while (node != null) {
            int cmp = compare(key,node.key) ;
            if (cmp == 0) {
                return node ;
            }
            if (cmp > 0) {
                //cmp > 0 说明element大，那么在右子树
                node = node.right ;
            } else {
                //cmp < 0
                node = node.left ;
            }
        }
        return null ;//说明没找到，那么返回null
    }
    private static class Node<K,V> {
        K key ;
        V value ;
        boolean color = RED;
        Node<K,V> left ;//左子节点
        Node<K,V> right ;//右子节点
        Node<K,V> parent ;//父节点
        int height ;
        public Node(K key, V value , Node<K,V> parent) {
            this.key = key ;
            this.value = value ;
            this.parent = parent ;//除了根节点，其余节点都具有父节点，所以要在构造方法传入
        }
        //判断是否为叶子节点
        public boolean isLeaf() {
            return left == null && right == null ;
        }
        //判断是否具有两个子节点
        public boolean hasTwoChildren() {
            return left != null && right != null ;
        }
        //说明是parent的左子树
        public boolean isLeftChild() {
            return parent != null && this == parent.left ;
        }
        //说明是parent的右子树
        public boolean isRightChild() {
            return parent != null && this == parent.right ;
        }
        /**
         * 找出兄弟节点
         * @return
         */
        public Node<K,V> sibling() {
            if (isLeftChild()) {
                //说明该节点是父节点的左子树
                return parent.right ;//说明该节点的兄弟节点即是父节点的右子树
            }
            if (isRightChild()) {
                //说明该节点是父节点的右子树
                return parent.left ;//说明该节点的兄弟节点即是父节点的左子树
            }
            //说明该节点既不是父节点的左子树 也不是父节点的右子树
            //那么说明该节点没有父节点，那么就没有兄弟节点
            return null ;
        }
    }
}