【190】Java8利用红黑树实现Map

红黑树的定义

红黑树是一种特殊类型的二叉搜索树（Binary Search Tree，缩写BST），除了满足二叉搜索树的定义外，还要满足下面五个红黑树特性：

1. 每个节点要么是红色，要么是黑色，必须二选一。
2. 根节点是黑色。
3. 每个叶子节点是黑色。叶子节点用空节点表示。
4. 红色节点的两个子节点都必须是黑色的。
5. 对于每个节点，从该节点到后代叶子节点的所有简单路径都包含相同数量的黑色节点。

红黑树的时间复杂度

对于N个节点的红黑树来说，它的搜索、插入、删除的时间复杂度都是 O(lgN) ，即以2为底N的对数。

与AVL树比较，红黑树的高度较高。红黑树不追求完全的平衡。所以红黑树搜索的性能略低于AVL树。

红黑树的插入和删除算法相对简单，红黑树的插入和删除性能高于AVL树。

代码实现

关于红黑树的代码放在 zhangchao.redblack 包里面。

RbTreeNode 类定义节点：

package zhangchao.redblack;

/**
 * 红黑树节点
 */
class RbTreeNode {
    // 键
    Object key;
    // 值
    Object value;
    // 左子节点
    RbTreeNode left = null;
    // 右子节点
    RbTreeNode right = null;
    // 父亲节点
    RbTreeNode p = null;
    // 颜色
    RbTreeColor color = null;
}

RbTreeColor是红黑树颜色的枚举

package zhangchao.redblack;

/**
 * 红黑树节点颜色的枚举
 */
enum RbTreeColor {
    RED, BLACK
}

RbTree 是红黑树类

package zhangchao.redblack;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Stack;

/**
 * 红黑树
 */
class RbTree {

    // 空节点
    static RbTreeNode nil;

    static {
        // 创建空节点并且设置成黑色
        nil = new RbTreeNode();
        nil.color = RbTreeColor.BLACK;
    }

    // 根节点
    RbTreeNode root = nil;

    private Comparator<Object> comparator;

    /**
     * 红黑树的构造方法
     * @param comparator 键的比较器
     */
    public RbTree(Comparator<Object> comparator) {
        this.comparator = comparator;
    }

    /**
     * 根据key获取节点
     * @param key 键
     * @return 节点
     */
    RbTreeNode getNode(Object key) {
        RbTreeNode current = this.root;
        while(nil != current) {
            int compareResult = this.comparator.compare(key, current.key);
            if (compareResult < 0) {
                current = current.left;
            } else if (compareResult > 0) {
                current = current.right;
            } else {
                return current;
            }
        }
        return nil;
    }


    /**
     * 生成节点列表。顺序是按照构造方法传入的比较器从小到大排序的。
     * @return 节点列表。顺序是按照构造方法传入的比较器从小到大排序的。
     */
    List<RbTreeNode> nodeList() {
        List<RbTreeNode> result = new ArrayList<>();
        if (nil == this.root) {
            return result;
        }
        RbTreeNode current = this.root;
        Stack<RbTreeNode> stack = new Stack<>();
        while(nil != current || !stack.isEmpty()) {
            while(nil != current) {
                stack.push(current);
                current = current.left;
            }
            current = stack.pop();
            // 放入结果列表中
            result.add(current);
            current = current.right;
        }
        return result;
    }

    /**
     * 插入新的节点z
     * @param z 新的节点
     */
    void insert(RbTreeNode z) {
        RbTreeNode x = this.root;
        RbTreeNode y = this.nil;
        while (x != nil) {
            y = x;
            if (this.comparator.compare(z.key, x.key) < 0) {
                x = x.left;
            } else {
                x = x.right;
            }
        }
        z.p = y;
        if (y == nil) {
            this.root = z;
        } else if (this.comparator.compare(z.key, y.key) < 0) {
            y.left = z;
        } else {
            y.right = z;
        }
        z.left = nil;
        z.right = nil;
        z.color = RbTreeColor.RED;
        this.insertFixup(z);
    }

    /**
     * 插入新的节点后，修复红黑树。避免出现不符合红黑树定义的情况。
     * @param z 新的节点
     */
    private void insertFixup(RbTreeNode z) {
        while (z.p.color == RbTreeColor.RED) {
            if (z.p == z.p.p.left) {
                RbTreeNode y = z.p.p.right; // y是z的叔叔
                if (y.color == RbTreeColor.RED) {
                    // case 1
                    // z的父亲和叔叔都是红色的。把父亲和叔叔都
                    // 染成黑色，爷爷染成红色，z指向爷爷节点。
                    z.p.color = RbTreeColor.BLACK;
                    y.color = RbTreeColor.BLACK;
                    z.p.p.color = RbTreeColor.RED;
                    z = z.p.p;
                } else {
                    if (z == z.p.right) {
                        // case 2
                        // z的叔叔是黑色，z是右子节点。对z执行左旋操作。
                        z = z.p;
                        this.leftRotate(z);
                    }
                    // case 3
                    // 可以承接 case 2
                    // 叔叔已经是黑色，再把父亲染成黑色，爷爷染成红色
                    // 然后右旋爷爷。
                    z.p.color = RbTreeColor.BLACK;
                    z.p.p.color = RbTreeColor.RED;
                    rightRotate(z.p.p);
                }
            } else {
                RbTreeNode y = z.p.p.left;
                if (y.color == RbTreeColor.RED) {
                    // case 1
                    // z的父亲和叔叔都是红色的。把父亲和叔叔都
                    // 染成黑色，爷爷染成红色，z指向爷爷节点。
                    z.p.color = RbTreeColor.BLACK;
                    y.color = RbTreeColor.BLACK;
                    z.p.p.color = RbTreeColor.RED;
                    z = z.p.p;
                } else {
                    // case 2
                    if (z == z.p.left) {
                        z = z.p;
                        this.rightRotate(z);
                    }
                    // case 3
                    z.p.color = RbTreeColor.BLACK;
                    z.p.p.color = RbTreeColor.RED;
                    this.leftRotate(z.p.p);
                }
            }
        }
        this.root.color = RbTreeColor.BLACK;
    }

    /**
     * 左旋。旋转的是x以及x的右子节点。
     * @param x 节点
     */
    private void leftRotate(final RbTreeNode x) {
        RbTreeNode y = x.right;
        x.right = y.left;
        if (y.left != nil) {
            y.left.p = x;
        }
        y.p = x.p;
        if (x.p == nil) {
            this.root = y;
        } else if (x == x.p.left) {
            x.p.left = y;
        } else {
            x.p.right = y;
        }
        y.left = x;
        x.p = y;
    }

    /**
     * 右旋。旋转的是x以及x的左子节点。
     * @param x 节点
     */
    private void rightRotate(final RbTreeNode x) {
        RbTreeNode y = x.left;
        x.left = y.right;
        if (y.right != nil) {
            y.right.p = x;
        }
        y.p = x.p;
        if (x.p == nil) {
            this.root = y;
        } else if (x == x.p.left) {
            x.p.left = y;
        } else {
            x.p.right = y;
        }
        y.right = x;
        x.p = y;
    }

    /**
     * 移植节点位置
     * @param u 要被删除的节点
     * @param v 替代u原来位置的节点
     */
    private void transplant(RbTreeNode u, RbTreeNode v) {
        if (u.p == nil) {
            this.root = v;
        } else if (u == u.p.left) {
            u.p.left = v;
        } else {
            u.p.right = v;
        }
        v.p = u.p;
    }

    /**
     * 以x为根节点的子树中，找到key最小的节点。
     * @param x 子树的根节点
     * @return 子树中key最小的节点
     */
    private RbTreeNode mininum(final RbTreeNode x) {
        RbTreeNode tmp = x;
        while (tmp.left != nil) {
            tmp = tmp.left;
        }
        return tmp;
    }

    /**
     * 删除节点
     * @param z 准备被删除的节点
     */
    void delete(RbTreeNode z) {
        RbTreeNode x = nil;
        RbTreeNode y = z;
        RbTreeColor yOriginalColor = y.color;
        if (z.left == nil) {
            x = z.right;
            transplant(z, z.right);
        } else if (z.right == nil) {
            x = z.left;
            transplant(z, z.left);
        } else {
            y = mininum(z.right);
            yOriginalColor = y.color;
            x = y.right;
            if (y != z.right) {
                transplant(y, y.right);
                y.right = z.right;
                y.right.p = y;
            } else {
                x.p = y;
            }
            transplant(z, y);
            y.left = z.left;
            y.left.p = y;
            y.color = z.color;
        }
        if (yOriginalColor == RbTreeColor.BLACK) {
            deleteFixup(x);
        }
    }

    /**
     * 删除节点后，修复红黑树，使得红黑树满足定义中的五个特性。
     * @param x 节点
     */
    private void deleteFixup(RbTreeNode x) {
        while (x != this.root && x.color == RbTreeColor.BLACK) {
            if (x == x.p.left) {
                RbTreeNode w = x.p.right; // w 是 x 的兄弟节点。
                if (w.color == RbTreeColor.RED) { // case 1
                    w.color = RbTreeColor.BLACK;
                    w.p.color = RbTreeColor.RED;
                    leftRotate(x.p);
                    w = x.p.right;
                }
                if (w.left.color == RbTreeColor.BLACK && w.right.color == RbTreeColor.BLACK) { // case 2
                    w.color = RbTreeColor.RED;
                    x = x.p;
                } else {
                    if (w.right.color == RbTreeColor.BLACK) { // case 3
                        w.left.color = RbTreeColor.BLACK;
                        w.color = RbTreeColor.RED;
                        rightRotate(w);
                        w = x.p.right;
                    }
                    // case 4
                    w.color = x.p.color;
                    x.p.color = RbTreeColor.BLACK;
                    w.right.color = RbTreeColor.BLACK;
                    leftRotate(x.p);
                    x = this.root;
                }
            } else {
                RbTreeNode w = x.p.left;
                if (w.color == RbTreeColor.RED) { // case 1
                    w.color = RbTreeColor.BLACK;
                    x.p.color = RbTreeColor.RED;
                    rightRotate(x.p);
                    w = x.p.left;
                }
                if (w.right.color == RbTreeColor.BLACK && w.left.color == RbTreeColor.BLACK) { // case 2
                    w.color = RbTreeColor.RED;
                    x = x.p;
                } else {
                    if (w.left.color == RbTreeColor.BLACK) { // case 3
                        w.right.color = RbTreeColor.BLACK;
                        w.color = RbTreeColor.RED;
                        leftRotate(w);
                        w = x.p.left;
                    }
                    // case 4
                    w.color = x.p.color;
                    x.p.color = RbTreeColor.BLACK;
                    w.left.color = RbTreeColor.BLACK;
                    rightRotate(x.p);
                    x = this.root;
                }
            }
        }
        x.color = RbTreeColor.BLACK;
    }
}

RedBlackTreeMap 基于红黑树实现的Map

package zhangchao.redblack;

import java.util.*;

/**
 * 基于红黑树实现的Map
 * @param  键
 * @param  值
 */
public class RedBlackTreeMap<K,V> implements Map<K,V>{

    // 红黑树对象
    private RbTree tree;

    // 键的比较器
    private Comparator<K> comparator;

    /**
     * 构造方法。
     * @param comparator 键的比较器
     */
    public RedBlackTreeMap(Comparator<K> comparator) {
        this.comparator = comparator;
        Comparator<Object> com = new Comparator<Object>() {
            @Override
            public int compare(Object o1, Object o2) {
                K k1 = (K)o1;
                K k2 = (K)o2;
                return comparator.compare(k1, k2);
            }
        };
        tree = new RbTree(com);
    }

    private void showTree(RbTreeNode node, int level, String prefix) {
        if (node == RbTree.nil) {
            return;
        }
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < level; i++) {
            sb.append("  ");
        }
        sb.append(prefix);
        sb.append(node.key).append("    ");
        if (node.color == RbTreeColor.RED) {
            sb.append("RED   ");
        } else if (node.color == RbTreeColor.BLACK) {
            sb.append("BLACK ");
        }
        if (node.p != RbTree.nil) {
            sb.append("p: ").append(node.p.key);
        }
        System.out.println(sb);
        level++;
        showTree(node.left, level, "left : ");
        showTree(node.right,level, "right: ");
    }

    public void showTree() {
        if (tree.root == RbTree.nil) {
            System.out.println("nil");
            return;
        }
        this.showTree(tree.root, 0, "root:");
    }

    /**
     * 包含键值的
     * @param  键
     * @param  值
     */
    static class RbEntry<K, V> implements Map.Entry<K, V> {
        K key;
        V value;

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public V setValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }
    }



    @Override
    public int size() {
        List<RbTreeNode> nodeList = tree.nodeList();
        return nodeList.size();
    }

    @Override
    public boolean isEmpty() {
        return (tree.root == RbTree.nil);
    }

    @Override
    public boolean containsKey(Object key) {
        RbTreeNode node = tree.getNode(key);
        return (tree.nil != node);
    }

    @Override
    public boolean containsValue(Object value) {
        List<RbTreeNode> nodeList = tree.nodeList();
        for (RbTreeNode item : nodeList) {
            if (item.value == null && null == value) {
                return true;
            }
            if (null != value && value.equals(item.value)) {
                return true;
            }
        }
        return false;
    }

    @Override
    public V get(Object key) {
        RbTreeNode node = tree.getNode((K)key);
        if (RbTree.nil != node) {
            return (V)node.value;
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        RbTreeNode node = tree.getNode(key);
        if (RbTree.nil != node) {
            V oldValue = (V)node.value;
            node.value = value;
            return oldValue;
        }
        RbTreeNode z = new RbTreeNode();
        z.key = key;
        z.value = value;
        z.left = RbTree.nil;
        z.right = RbTree.nil;
        z.p = RbTree.nil;
        z.color = RbTreeColor.RED;
        tree.insert(z);
        return null;
    }

    @Override
    public V remove(Object key) {
        RbTreeNode z = tree.getNode(key);
        if (z != RbTree.nil) {
            V oldValue = (V)z.value;
            tree.delete(z);
            return oldValue;
        }
        return null;
    }

    @Override
    public void putAll(Map<? extends K, ? extends V> m) {
        if (null == m || m.isEmpty()) {
            return;
        }
        Set<? extends K> keySet = m.keySet();
        for (K k : keySet) {
            V v = m.get(k);
            this.put(k, v);
        }
    }

    @Override
    public void clear() {
        tree.root = RbTree.nil;
    }

    @Override
    public Set<K> keySet() {
        List<RbTreeNode> nodeList = tree.nodeList();
        Set<K> set = new TreeSet<K>(this.comparator);
        for (RbTreeNode item : nodeList) {
            set.add((K)item.key);
        }
        return set;
    }

    @Override
    public Collection<V> values() {
        List<V> result = new ArrayList<>();
        List<RbTreeNode> nodeList = tree.nodeList();
        for (RbTreeNode item : nodeList) {
            result.add((V)item.value);
        }
        return result;
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        List<RbTreeNode> nodeList = tree.nodeList();
        List<RbEntry<K, V>> list = new ArrayList<>();
        for (RbTreeNode item : nodeList) {
            RbEntry<K, V> rbEntry = new RbEntry<>();
            rbEntry.key = (K)item.key;
            rbEntry.value = (V)item.value;
            list.add(rbEntry);
        }
        TreeSet<Entry<K, V>> set = new TreeSet<>(new Comparator<Entry<K, V>>() {
            @Override
            public int compare(Entry<K, V> o1, Entry<K, V> o2) {
                K k1 = o1.getKey();
                K k2 = o2.getKey();
                return comparator.compare(k1, k2);
            }
        });
        for (RbEntry<K, V> item : list) {
            set.add(item);
        }
        return set;
    }
}

Test1是测试类

package zhangchao.redblack.test;


import zhangchao.avl.AVLTreeMap;
import zhangchao.redblack.RedBlackTreeMap;

import java.util.*;

public class Test1 {
    public static void t1() {
        Map<Integer, String> map = new RedBlackTreeMap<>( (o1, o2) ->{
            if (null == o1 && null == o2) {
                return 0;
            }
            if (null == o1 && null != o2) {
                return -1;
            }
            if (null != o1 && null == o2) {
                return 1;
            }
            return o1 - o2;
        });
        int[] arr = new int[]{8,3,6,1,2,98,2,6,150,170,160,7,52,28,75,14,
                40,86,10,21,46,25};

        for (int i : arr) {
            map.put(i, "__" + String.valueOf(i));
        }

        RedBlackTreeMap redBlackTreeMap = (RedBlackTreeMap) map;
        redBlackTreeMap.showTree();

        System.out.println(map.get(3));
        System.out.println(map.get(6));
        System.out.println(map.get(98));
        System.out.println(map.get(null));

        Set<Integer> set = redBlackTreeMap.keySet();
        for (Integer i : set) {
            System.out.println(i);
        }
        System.out.println();

        HashSet<Integer> hashSet = new HashSet<>();
        for (int i : arr) {
            hashSet.add(i);
        }
        for (int i : hashSet) {
            if (!set.contains(i)) {
                System.out.println(false);
            }
        }
        System.out.println(set.size() + "   " + hashSet.size());

        System.out.println("containsKey 3: " + redBlackTreeMap.containsKey(3));
        System.out.println("containsKey 4: " + redBlackTreeMap.containsKey(4));
        System.out.println("containsValue __3: " + redBlackTreeMap.containsValue("__3"));
        System.out.println("containsValue __4: " + redBlackTreeMap.containsValue("__4"));
        System.out.println();

        Set<Map.Entry<Integer, String>> entrySet = redBlackTreeMap.entrySet();
        for (Map.Entry<Integer, String> item : entrySet) {
            System.out.println(item.getKey() + ": " + item.getValue());
        }
//        avlTreeMap.checkTree();
    }


    public static void t2() {
        Map<Integer, String> map = new RedBlackTreeMap<>((o1, o2) ->{
            if (null == o1 && null == o2) {
                return 0;
            }
            if (null == o1 && null != o2) {
                return -1;
            }
            if (null != o1 && null == o2) {
                return 1;
            }
            return o1 - o2;
        });
        int[] arr = new int[]{8,3,6,1,2,98,2,6,150,170,160,7,52,28,75,14,
                40,86,10,21,46,25};
        for (int i : arr) {
            map.put(i, "__" + String.valueOf(i));
        }
        RedBlackTreeMap redBlackTreeMap = (RedBlackTreeMap) map;
        redBlackTreeMap.showTree();
        redBlackTreeMap.remove(7);
        System.out.println("\n\n\n");
        redBlackTreeMap.showTree();
//        redBlackTreeMap.checkTree();
        System.out.println("get(7):  " + map.get(7));
    }

    public static void t3_get(Map<Integer, String> map, List<Integer> list, String label) {
        long t1 = System.currentTimeMillis();
        for (Integer item : list) {
            map.get(item);
        }
        long t2 = System.currentTimeMillis();
        long t3 = t2 - t1;
        System.out.println(label + " time: " + t3);
    }

    public static void t3_remove(Map<Integer, String> map, List<Integer> list, String label) {
        long t1 = System.currentTimeMillis();
        for (Integer item : list) {
            map.remove(item);
        }
        long t2 = System.currentTimeMillis();
        long t3 = t2 - t1;
        System.out.println(label + " time: " + t3);
    }

    public static void t3_insert(Map<Integer, String> map, List<Integer> list, String label) {
        long t1 = System.currentTimeMillis();
        for (Integer item : list) {
            map.put(item, "__" + item);
        }
        long t2 = System.currentTimeMillis();
        long t3 = t2 - t1;
        System.out.println(label + " time: " + t3);
    }

    public static void t3() {
        Comparator<Integer> comparator = (o1, o2) ->{
            if (null == o1 && null == o2) {
                return 0;
            }
            if (null == o1 && null != o2) {
                return -1;
            }
            if (null != o1 && null == o2) {
                return 1;
            }
            return o1 - o2;
        };
        RedBlackTreeMap<Integer, String> redBlackTreeMap = new RedBlackTreeMap<>(comparator);
        AVLTreeMap<Integer, String> avlTreeMap = new AVLTreeMap(comparator);
        final int a = 100000;
        List<Integer> list = new ArrayList<>();
        Random random = new Random();
        for (int i = 0; i < a; i++) {
            list.add(random.nextInt(a));
        }

        t3_insert(redBlackTreeMap, list, "red insert");
        t3_insert(avlTreeMap, list, "avl insert");
        System.out.println();
        Collections.shuffle(list);
        t3_get(redBlackTreeMap, list, "red get");
        t3_get(avlTreeMap, list, "avl get");
        System.out.println();
        t3_remove(redBlackTreeMap, list, "red remove");
        t3_remove(avlTreeMap, list, "avl remove");
    }

    public static void t4() {
        Comparator<Integer> comparator = (o1, o2) ->{
            if (null == o1 && null == o2) {
                return 0;
            }
            if (null == o1 && null != o2) {
                return -1;
            }
            if (null != o1 && null == o2) {
                return 1;
            }
            return o1 - o2;
        };
        RedBlackTreeMap<Integer, String> redBlackTreeMap = new RedBlackTreeMap<>(comparator);
        for (int i = 0; i < 10; i++) {
            redBlackTreeMap.put(i, "__" + i);
        }
        Set<Map.Entry<Integer, String>> entrySet = redBlackTreeMap.entrySet();
        for (Map.Entry<Integer, String> item : entrySet) {
            System.out.println(item.getKey() + "   " + item.getValue());
        }
    }

    public static void main(String[] args) {
        t4();

    }
}