中招了，重写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>() {@Overridepublic int compare(Long o1, Long o2) {// 1、如果v1等于v2，则值为0// 2、如果v1小于v2，则值为-1// 3、如果v1等于v2，则值为1Integer 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、基于排序数据的线性时间树构建算法，进行buildbuildFromSorted(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;// 修改次数加1modCount++;return null;}int cmp;TreeMap.Entry<K,V> parent;// 获取比较器Comparator<? super K> cpr = comparator;// 三、如果比较器不为空，则执行一下逻辑，即自定义比较器执行逻辑if (cpr != null) {do {// 1、将t节点赋值给parentparent = 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;elseparent.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，则值为1Integer 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;elsereturn 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，则将左节点的值赋值给pif (cmp < 0)p = p.left;// 3、如果大于0，则将右节点的值赋值给pelse if (cmp > 0)p = p.right;else// 4、如果等于0，则返回p节点return p;}}return null;
}

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

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

问题找到了，解决吧！