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();
        }
    }
}
最近下载更多
matintalorr  LV10 2021年8月31日
最代码官方  LV167 2020年1月11日
最近浏览更多
graceful 2023年10月18日
暂无贡献等级
666ing  LV2 2023年2月26日
学不会代码就去种地 2022年11月22日
暂无贡献等级
绘飛的渔 2022年8月2日
暂无贡献等级
yuyan701  LV1 2022年1月1日
matintalorr  LV10 2021年8月31日
smile 2021年6月10日
暂无贡献等级
絮落无痕  LV13 2021年4月21日
Killah  LV9 2021年4月16日
zarathurstra 2020年11月6日
暂无贡献等级
顶部 客服 微信二维码 底部
>扫描二维码关注最代码为好友扫描二维码关注最代码为好友