package com.mjy.map;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Stack;

/**
 * ***************************************************
 * @ClassName TreeMap
 * @Description 描述
 * @Author maojianyun
 * @Date 2020/1/9 22:06
 * @Version V1.0
 * ****************************************************
 **/
public class TreeMap<K, V> implements Map<K, V>{

    private final static boolean RED = false;

    private final static boolean BLACK = true;

    private Node<K, V> root;

    private int size;

    private Comparator<K> comparator;

    public TreeMap() {
        this(null);
    }

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

    public int size() {
        return size;
    }

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

    public void clear() {
        root = null;
    }

    public V put(K key, V value) {
        keyNotNullCheck(key);
        if (root == null) { // 添加根节点
            Node<K, V> node = new Node<K, V>(key, value, null);
            root = node;
            size++;
            afterPut(node);
            return null;
        }
        Node<K, V> temp = root;
        Node<K, V> parent = root;
        int com = 0;
        while (temp != null){
            parent = temp;
            com = compare(key, temp.key);
            if (com > 0) {
                temp = temp.right;
            } else if (com < 0){
                temp = temp.left;
            } else { // 相等、覆盖原来的值并且返回
                temp.key = key;
                V oldValue = temp.value;
                temp.value = value;
                return oldValue;
            }
        }
        Node<K, V> newNode = new Node<K, V>(key, value, parent);
        if (com > 0) {
            parent.right = newNode;
        } else {
            parent.left = newNode;
        }
        size++;
        afterPut(newNode);
        return null;
    }

    public V get(K key) {
        Node<K, V> node = node(key);
        return node != null? node.value: null;
    }

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

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

    public boolean containsValue(V value) {
        if (value == null) {
            return false;
        }
        Queue<Node<K, V>> queue = new LinkedList<Node<K, V>>();
        queue.offer(root);
        while(queue.size() > 0){
            Node<K, V> node = queue.poll();
            // 比较值
            if (valEquals(value, node.value)) {
                return true;
            }
            if (node.left != null) {
                queue.offer(node.left);
            }
            if (node.right != null) {
                queue.offer(node.right);
            }
        }
        return false;
    }


    public void traversal(Visitor<K, V> visitor) {
        if (root == null) {
            return;
        }
        Node<K, V> temp = root;
        Stack<Node<K, V>> stack = new Stack<Node<K, V>>();
        while (temp != null || !stack.isEmpty()) {
            if (temp != null) {
                stack.push(temp);
                temp = temp.left;
            } else {
                if (visitor.stop) {
                    break;
                }
                Node<K, V> node = stack.pop();
                visitor.stop = visitor.visit(node.key, node.value);
                temp = node.right;
            }
        }
    }

    private boolean valEquals(V v1, V v2) {
        return v1 == null ? v2 == null : v1.equals(v2);
    }

    /**
     * 删除node
     * @param node
     * @return
     */
    private V remove(Node<K, V> node){
        if (node == null) {
            return null;
        }
        V oldValue = node.value;
        if (node.hasTwoChildren()){ // 度为2的节点
            // 找到后继节点
            Node<K, V> su = successor(node);
            // 用后继节点的值覆盖度为2的节点的值
            node.key = su.key;
            node.value = su.value;
            // 删除后继节点
            node = su;
        }
        // 删除node节点(node的度必然是1或者0)
        // 如果度为1、则返回的是其子节点
        Node replacement = node.left != null? node.left: node.right;
        if (replacement != null) { // 删除度为1的节点
            replacement.parent = node.parent;
            if (node.parent == null) {
                root = replacement;
            }else if (node.isLeftChild()) {
                node.parent.left = replacement;
            } if (node.isRightChild()) {
                node.parent.right = replacement;
            }
            // 删除节点之后的处理
            afterRemove(replacement);
        } else if (node.parent == null) { // node是叶子节点并且是根节点
           root = null;
        } else { // node是叶子节点,但不是根节点
            if (node == node.parent.left) {
                node.parent.left = null;
            } else { // node == node.parent.right
                node.parent.right = null;
            }
            // 删除节点之后的处理
            afterRemove(node);
        }
        size--;
        return oldValue;
    }
    /**
     * 根据可以得到node
     * @param key
     * @return
     */
    private Node<K, V> node(K key) {
        keyNotNullCheck(key);
        Node<K, V> temp = root;
        while (temp != null) {
            int com = compare(key, temp.key);
            if (com == 0) {
                return  temp;
            } else if (com > 0) {
                temp = temp.right;
            } else {
                temp = temp.left;
            }
        }
        return null;
    }

    /**
     * 得到节点的后继节点
     * @param node
     * @return
     */
    private Node<K, V> successor(Node<K, V> node){
        if (node == null) return null;

        // 前驱节点在左子树当中(right.left.left.left....)
        Node<K, V> p = node.right;
        if (p != null) {
            while (p.left != null) {
                p = p.left;
            }
            return p;
        }

        // 从父节点、祖父节点中寻找前驱节点
        while (node.parent != null && node == node.parent.right) {
            node = node.parent;
        }
        return node.parent;
    }

    /**
     * key的比较:返回值大于0说明key1 > key2, 返回值等于0说明key1 == key2 ,返回值小于0说key1 < key2
     * @param key1
     * @param key2
     * @return
     */
    private int compare(K key1, K key2){
        if (comparator != null) {
            return comparator.compare(key1, key2);
        }
        return ((Comparable)key1).compareTo(key2);
    }

    /**
     * 校验可以是否为空
     * @param key
     */
    private void keyNotNullCheck(K key){
        if (key == null) {
            throw new IllegalArgumentException("key must not be null");
        }
    }

    /**
     * 给指定节点染指定的颜色
     * @param node 节点
     * @param color 颜色
     * @return 节点
     */
    private Node<K, V> color(Node<K, V> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
        return node;
    }

    /**
     * 节点染成红色
     * @param node
     * @return
     */
    private Node<K, V> red(Node<K, V> node){
        return color(node, RED);
    }

    /**
     * 节点染成黑色
     * @param node
     * @return
     */
    private Node<K, V> black(Node<K, V> node){
        return  color(node, BLACK);
    }

    /**
     * 得到节点的颜色
     * @param node
     * @return
     */
    private boolean colorOf(Node<K, V> node){
        return node == null? BLACK: node.color;
    }

    /**
     * 是否是黑色
     * @param node
     * @return
     */
    private boolean isBlack(Node<K, V> node) {
        return colorOf(node) == BLACK;
    }

    /**
     * 是否是红色
     * @param node
     * @return
     */
    private boolean isRed(Node<K, V> node) {
        return colorOf(node) == RED;
    }

    private void afterPut(Node<K, V> node) {
        Node<K, V> parent = node.parent;
        // 添加的是根节点 或者 上溢到达了根节点
        if (parent == null) {
            black(node);
            return;
        }
        // 如果父节点是黑色,直接返回
        if (isBlack(parent)) {
            return;
        }
        // 叔父节点
        Node<K, V> uncle = parent.sibling();
        // 祖父节点
        Node<K, V> grand = red(parent.parent);
        if (isRed(uncle)) { // 叔父节点是红色【B树节点上溢】
            black(parent);
            black(uncle);
            // 把祖父节点当做是新添加的节点
            afterPut(grand);
            return;
        }
        // 叔父节点不是红色
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                black(parent);
            } else { // LR
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        } else { // R
            if (node.isLeftChild()) { // RL
                black(node);
                rotateRight(parent);
            } else { // RR
                black(parent);
            }
            rotateLeft(grand);
        }
    }

    private void afterRemove(Node<K, V> node) {
        // 如果删除的节点是红色
        // 或者 用以取代删除节点的子节点是红色
        if (isRed(node)) {
            black(node);
            return;
        }
        Node<K, V> parent = node.parent;
        if (parent == null) {
            return;
        }
        // 删除的是黑色叶子节点【下溢】
        // 判断被删除的node是左还是右
        boolean left = parent.left == null || node.isLeftChild();
        Node<K, V> sibling = left ? parent.right : parent.left;
        if (left) { // 被删除的节点在左边,兄弟节点在右边
            if (isRed(sibling)) { // 兄弟节点是红色
                black(sibling);
                red(parent);
                rotateLeft(parent);
                // 更换兄弟
                sibling = parent.right;
            }

            // 兄弟节点必然是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点没有1个红色子节点,父节点要向下跟兄弟节点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else { // 兄弟节点至少有1个红色子节点,向兄弟节点借元素
                // 兄弟节点的左边是黑色,兄弟要先旋转
                if (isBlack(sibling.right)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }

                color(sibling, colorOf(parent));
                black(sibling.right);
                black(parent);
                rotateLeft(parent);
            }
        } else { // 被删除的节点在右边,兄弟节点在左边
            if (isRed(sibling)) { // 兄弟节点是红色
                black(sibling);
                red(parent);
                rotateRight(parent);
                // 更换兄弟
                sibling = parent.left;
            }

            // 兄弟节点必然是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点没有1个红色子节点,父节点要向下跟兄弟节点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else { // 兄弟节点至少有1个红色子节点,向兄弟节点借元素
                // 兄弟节点的左边是黑色,兄弟要先旋转
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling);
                    sibling = parent.left;
                }

                color(sibling, colorOf(parent));
                black(sibling.left);
                black(parent);
                rotateRight(parent);
            }
        }
    }

    private void rotateLeft(Node<K, V> grand) {
        Node<K, V> parent = grand.right;
        Node<K, V> child = parent.left;
        grand.right = child;
        parent.left = grand;
        afterRotate(grand, parent, child);
    }

    private void rotateRight(Node<K, V> grand) {
        Node<K, V> parent = grand.left;
        Node<K, V> child = parent.right;
        grand.left = child;
        parent.right = grand;
        afterRotate(grand, parent, child);
    }

    private void afterRotate(Node<K, V> grand, Node<K, V> parent, Node<K, V> child) {
        // 让parent称为子树的根节点
        parent.parent = grand.parent;
        if (grand.isLeftChild()) {
            grand.parent.left = parent;
        } else if (grand.isRightChild()) {
            grand.parent.right = parent;
        } else { // grand是root节点
            root = parent;
        }

        // 更新child的parent
        if (child != null) {
            child.parent = grand;
        }

        // 更新grand的parent
        grand.parent = parent;
    }

    /**
     * ***************************************************
     * @ClassName Node<K, V>
     * @Description 红黑树节点
     * @Author maojianyun
     * @Date 2020/1/9 22:06
     * @Version V1.0
     * ****************************************************
     **/
    private static class Node<K, V>{
        // 默认为红色
        boolean color = RED;
        K key;
        V value;
        Node<K, V> left;
        Node<K, V> right;
        Node<K, V> parent;
        public Node(K key, V value, Node<K, V> parent){
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        /**
         * 是否是叶子节点
         * @return
         */
        public boolean isLeaf(){
            return left == null && right == null;
        }

        /**
         * 是否有两个孩子
         * @return
         */
        public boolean hasTwoChildren(){
            return left != null && right != null;
        }

        /**
         * 是否是左孩子
         * @return
         */
        public boolean isLeftChild(){
            return parent != null && this == parent.left;
        }

        /**
         * 是否是右孩子
         * @return
         */
        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;
        }

        @Override
        public String toString() {
            String str = "";
            if (color == RED) {
                str = "R_";
            }
            return str + key.toString() + "_" + value.toString();
        }
    }
}
最近下载更多
最代码官方 LV1582020年1月11日
皇冠皇冠太阳月亮月亮月亮星星星星
最近浏览更多
zarathurstra2020年11月6日
暂无贡献等级
xwq1234567 LV12020年10月28日
星星
pxqtsht LV132020年7月3日
月亮月亮月亮星星
11111155255252020年6月28日
暂无贡献等级
17600446733 LV202020年6月23日
太阳月亮
guadan2020年5月30日
暂无贡献等级
Jacko01 LV62020年5月18日
月亮星星星星
加油干阳神 LV52020年5月13日
月亮星星
seychell2020年3月22日
暂无贡献等级
bsszds2332020年3月10日
暂无贡献等级
顶部客服微信二维码底部
>扫描二维码关注最代码为好友扫描二维码关注最代码为好友