中招了，重写TreeMap的比较器引发的问题…

需求背景

• 给一个无序的map，按照value的值进行排序，value值越小，排在越前面。
• key和value都不为null
• value可能相同
• 返回结果为一个相同的有序map

代码如下所示:

// 假设，key=商品id，value=商品剩余库存
Map<Long, Integer> map = new HashMap<>();
map.put(1L, 10);
map.put(2L, 20);
map.put(3L, 10);

到这里，大家可以先想想，如果是你会怎么解决？

我的解决思路

1、使用TreeMap，因为TreeMap可以对元素进行排序

2、重写TreeMap的比较器

代码如下所示：

// 承接上面的代码
// 按照 value 排序
Map<Long, Integer> treeMap1 = new TreeMap<>(new Comparator<Long>() {
    @Override
    public int compare(Long o1, Long o2) {
        // 1、如果v1等于v2，则值为0
        // 2、如果v1小于v2，则值为-1
        // 3、如果v1等于v2，则值为1
        Integer value1 = map.get(o1);
        Integer value2 = map.get(o2);
        return value1.compareTo(value2);
    }
});
treeMap1.putAll(map);
System.out.println(treeMap1);

运行后的结果为：

{1=10, 2=20}

what？为什么我们添加了3个元素，结果少了一个呢？

TreeMap putAll源码分析

让我们来看看 putAll 的具体过程

1、分析 TreeMap.putAll

源码如下所示:

public void putAll(Map<? extends K, ? extends V> map) {
    // 一、获取待添加的map的大小
    int mapSize = map.size();
    // 二、当前的size大小等于0 且 待添加的map的大小不等于0 且 待添加的map是SortedMap的实现类，则执行以下逻辑
    if (size==0 && mapSize!=0 && map instanceof SortedMap) {
        // 1、获取待添加的map的比较器
        Comparator<?> c = ((SortedMap<?,?>)map).comparator();
        // 2、如果两个比较器相同，则执行以下逻辑
        if (c == comparator || (c != null && c.equals(comparator))) {
            // 3、修改次数+1
            ++modCount;
            try {
                // 4、基于排序数据的线性时间树构建算法，进行build
                buildFromSorted(mapSize, map.entrySet().iterator(),
                        null, null);
            } catch (java.io.IOException cannotHappen) {
            } catch (ClassNotFoundException cannotHappen) {
            }
            return;
        }
    }
    // 三、如果不符合上面的条件，则执行父类的 putAll 方法
    super.putAll(map);
}

从上面源码，不难看出，我们的数据符合 流程二，但是不符合 流程二-2，所以我们会执行父类的 putAll 方法，即流程三。

2、分析 AbstractMap.putAll

TreeMap 继承 AbstractMap，所以 super.putAll(map)，执行的 putAll 为 AbstractMap 的 putAll 方法，源码如下所示：

public void putAll(Map<? extends K, ? extends V> m) {
    // 遍历 m map，将它所有的值，使用put方法，全部添加到当前的map中
    for (Map.Entry<? extends K, ? extends V> e : m.entrySet())
        put(e.getKey(), e.getValue());
}

这段代码简单，就是一个遍历添加元素的。

但是有一个问题，这里的 put 方法执行的是谁的 put 方法呢？

• 1、AbstractMap.put
• 2、TreeMap.put

这里大家可以先思考1分钟，然后再继续往下看。

答案是：

执行的是 TreeMap.put

回答错误 or 不知道真实原因的小伙伴，可以去网上搜搜答案，这里是一个很重要的基础知识点哦。

3、分析 TreeMap.put

源代码如下所示：

public V put(K key, V value) {
    // 一、获取根节点
    TreeMap.Entry<K,V> t = root;
    // 二、判断跟节点是否为空
    if (t == null) {
        // 类型检查 and null 检查
        compare(key, key); // type (and possibly null) check
        // 创建根节点
        root = new TreeMap.Entry<>(key, value, null);
        size = 1;
        // 修改次数加1
        modCount++;
        return null;
    }

    int cmp;
    TreeMap.Entry<K,V> parent;
    // 获取比较器
    Comparator<? super K> cpr = comparator;
    // 三、如果比较器不为空，则执行一下逻辑，即自定义比较器执行逻辑
    if (cpr != null) {
        do {
            // 1、将t节点赋值给parent
            parent = t;
            // 2、比较t节点的key是否与待添加的key相等
            cmp = cpr.compare(key, t.key);
            // 3、如果返回值小于0，则将左子树赋值给t节点，即后续遍历左子树
            if (cmp < 0)
                t = t.left;
            // 4、如果返回值大于0，则将右子树赋值给t节点，即后续遍历右子树
            else if (cmp > 0)
                t = t.right;
            else
                // 5、如果返回值为0，则覆盖原来的值
                return t.setValue(value);
        } while (t != null);
    }
    // 四、如果比较器为空，则执行以下逻辑，即默认执行逻辑
    else {
      // 这部分逻辑，先忽略
    }
    TreeMap.Entry<K,V> e = new TreeMap.Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

我们结合上面的源码和我们自定义的排序器，就可以发现以下问题：

1、我们比较的是两个 value 的大小，而 value 可能是一样的。

这种情况下，就会覆盖原来的值，这个就是我们执行 putAll 后，元素缺失的原因了。

好了既然问题找到了，那如何解决这个问题呢？

如果是你，你会怎么解决呢？可以花一分钟时间思考一下，再看后面的内容。

4、解决 TreeMap.putAll，元素缺失的问题

我当时想到最直接的方案就是，在 value 相等的情况下，不返回 0，返回1 or -1，这样就可以最简单、最快捷的解决这个问题了。

修改后的代码如下所示：

// 这里换了一种写法，是java8的特性，简化了代码（为了偷懒）
Map<Long, Integer> treeMap2 = new TreeMap<>((key1, key2) -> {
    // 1、如果v1等于v2，则值为0
    // 2、如果v1小于v2，则值为-1
    // 3、如果v1等于v2，则值为1
    Integer value1 = map.get(key1);
    Integer value2 = map.get(key2);

    int result = value1.compareTo(value2);

    if (result == 0) {
        return -1;
    }
    return result;
});

treeMap2.putAll(map);
System.out.println(treeMap2);

运行后的结果为：

{3=10, 1=10, 2=20}

我们可以发现，3个值都有了，并且是有序的，完美符合需求！好了，关机下班！

然而事情并没有结束 （大家可以想一下，这样写会有什么问题呢？）！

新的问题出现

第二天，高高兴兴的写着业务代码、调试逻辑，突然一个 空指针 的报错，出现了。这也太常见了吧，3分钟内解决！

排查了半天，发现又回到了昨天的修改的那段逻辑了。

1、TreeMap.get 获取不到值

简化版代码如下所示：

// 假设，key=商品id，value=商品剩余库存
Map<Long, Integer> map = new HashMap<>();
map.put(1L, 10);
map.put(2L, 20);
map.put(3L, 10);

// 排序
Map<Long, Integer> treeMap2 = new TreeMap<>((key1, key2) -> {
    Integer value1 = map.get(key1);
    Integer value2 = map.get(key2);

    int result = value1.compareTo(value2);

    if (result == 0) {
        return -1;
    }
    return result;
});
treeMap2.putAll(map);
System.out.println(treeMap2);

// 获取商品1的剩余数量
Integer quantity = treeMap2.get(1L);
System.out.println(quantity);

运行后的结果为：

{3=10, 1=10, 2=20}
null

这个结果令我百思不得其解，只能看看源码咯。

2、分析 TreeMap.get

源码如下所示：

public V get(Object key) {
    // 根据key获取节点
    TreeMap.Entry<K,V> p = getEntry(key);
    // 节点为空则返回null，否则返回节点的 value 值
    return (p==null ? null : p.value);
}

final TreeMap.Entry<K,V> getEntry(Object key) {
    // 一、如果比较器不为空，则执行一下逻辑
    if (comparator != null)
        // 1、使用自定义比较器取出key对应的节点
        return getEntryUsingComparator(key);
    // 二、如果比较器为空，且key为null，则抛空指针异常
    if (key == null)
        throw new NullPointerException();
    @SuppressWarnings("unchecked")
    Comparable<? super K> k = (Comparable<? super K>) key;
    TreeMap.Entry<K,V> p = root;
    // 三、取出key对应的节点
    while (p != null) {
        int cmp = k.compareTo(p.key);
        if (cmp < 0)
            p = p.left;
        else if (cmp > 0)
            p = p.right;
        else
            return p;
    }
    return null;
}

从上面的源码，我们可以发现，问题肯定就是出现在 getEntryUsingComparator 方法里了。

2、分析 TreeMap.getEntryUsingComparator

源码如下所示：

final TreeMap.Entry<K,V> getEntryUsingComparator(Object key) {
    // 一、将key转换成对应的类型
    @SuppressWarnings("unchecked")
    K k = (K) key;
    // 二、获取比较器
    Comparator<? super K> cpr = comparator;
    // 三、判断比较器是否为空
    if (cpr != null) {
        // 1、遍历map，取出key对应的节点对象
        TreeMap.Entry<K,V> p = root;
        while (p != null) {
            int cmp = cpr.compare(k, p.key);
            // 2、如果小于0，则将左节点的值赋值给p
            if (cmp < 0)
                p = p.left;
            // 3、如果大于0，则将右节点的值赋值给p
            else if (cmp > 0)
                p = p.right;
            else
                // 4、如果等于0，则返回p节点
                return p;
        }
    }
    return null;
}

结合上面的源码，和我们之前自定义的比较器，我们不难发现问题出现在哪里：

自定义比较器，没有返回0的情况

问题找到了，解决吧！