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

需求背景

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

代码如下所示:

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

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

我的解决思路

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

2、重写TreeMap的比较器

代码如下所示：

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

运行后的结果为：

{1=10, 2=20}

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

TreeMap putAll源码分析

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

1、分析 TreeMap.putAll

源码如下所示:

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

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

2、分析 AbstractMap.putAll

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

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

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

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

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

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

答案是：

执行的是 TreeMap.put

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

3、分析 TreeMap.put

源代码如下所示：

 1 public V put(K key, V value) {
 2     // 一、获取根节点
 3     TreeMap.Entry<K,V> t = root;
 4     // 二、判断跟节点是否为空
 5     if (t == null) {
 6         // 类型检查 and null 检查
 7         compare(key, key); // type (and possibly null) check
 8         // 创建根节点
 9         root = new TreeMap.Entry<>(key, value, null);
10         size = 1;
11         // 修改次数加1
12         modCount++;
13         return null;
14     }
15
16     int cmp;
17     TreeMap.Entry<K,V> parent;
18     // 获取比较器
19     Comparator<? super K> cpr = comparator;
20     // 三、如果比较器不为空，则执行一下逻辑，即自定义比较器执行逻辑
21     if (cpr != null) {
22         do {
23             // 1、将t节点赋值给parent
24             parent = t;
25             // 2、比较t节点的key是否与待添加的key相等
26             cmp = cpr.compare(key, t.key);
27             // 3、如果返回值小于0，则将左子树赋值给t节点，即后续遍历左子树
28             if (cmp < 0)
29                 t = t.left;
30             // 4、如果返回值大于0，则将右子树赋值给t节点，即后续遍历右子树
31             else if (cmp > 0)
32                 t = t.right;
33             else
34                 // 5、如果返回值为0，则覆盖原来的值
35                 return t.setValue(value);
36         } while (t != null);
37     }
38     // 四、如果比较器为空，则执行以下逻辑，即默认执行逻辑
39     else {
40       // 这部分逻辑，先忽略
41     }
42     TreeMap.Entry<K,V> e = new TreeMap.Entry<>(key, value, parent);
43     if (cmp < 0)
44         parent.left = e;
45     else
46         parent.right = e;
47     fixAfterInsertion(e);
48     size++;
49     modCount++;
50     return null;
51 }

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

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

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

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

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

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

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

修改后的代码如下所示：

 1 // 这里换了一种写法，是java8的特性，简化了代码（为了偷懒）
 2 Map<Long, Integer> treeMap2 = new TreeMap<>((key1, key2) -> {
 3     // 1、如果v1等于v2，则值为0
 4     // 2、如果v1小于v2，则值为-1
 5     // 3、如果v1等于v2，则值为1
 6     Integer value1 = map.get(key1);
 7     Integer value2 = map.get(key2);
 8
 9     int result = value1.compareTo(value2);
10
11     if (result == 0) {
12         return -1;
13     }
14     return result;
15 });
16
17 treeMap2.putAll(map);
18 System.out.println(treeMap2);

运行后的结果为：

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

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

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

新的问题出现

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

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

1、TreeMap.get 获取不到值

简化版代码如下所示：

 1 // 假设，key=商品id，value=商品剩余库存
 2 Map<Long, Integer> map = new HashMap<>();
 3 map.put(1L, 10);
 4 map.put(2L, 20);
 5 map.put(3L, 10);
 6
 7 // 排序
 8 Map<Long, Integer> treeMap2 = new TreeMap<>((key1, key2) -> {
 9     Integer value1 = map.get(key1);
10     Integer value2 = map.get(key2);
11
12     int result = value1.compareTo(value2);
13
14     if (result == 0) {
15         return -1;
16     }
17     return result;
18 });
19 treeMap2.putAll(map);
20 System.out.println(treeMap2);
21
22 // 获取商品1的剩余数量
23 Integer quantity = treeMap2.get(1L);
24 System.out.println(quantity);

运行后的结果为：

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

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

2、分析 TreeMap.get

源码如下所示：

 1 public V get(Object key) {
 2     // 根据key获取节点
 3     TreeMap.Entry<K,V> p = getEntry(key);
 4     // 节点为空则返回null，否则返回节点的 value 值
 5     return (p==null ? null : p.value);
 6 }
 7
 8 final TreeMap.Entry<K,V> getEntry(Object key) {
 9     // 一、如果比较器不为空，则执行一下逻辑
10     if (comparator != null)
11         // 1、使用自定义比较器取出key对应的节点
12         return getEntryUsingComparator(key);
13     // 二、如果比较器为空，且key为null，则抛空指针异常
14     if (key == null)
15         throw new NullPointerException();
16     @SuppressWarnings("unchecked")
17     Comparable<? super K> k = (Comparable<? super K>) key;
18     TreeMap.Entry<K,V> p = root;
19     // 三、取出key对应的节点
20     while (p != null) {
21         int cmp = k.compareTo(p.key);
22         if (cmp < 0)
23             p = p.left;
24         else if (cmp > 0)
25             p = p.right;
26         else
27             return p;
28     }
29     return null;
30 }

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

2、分析 TreeMap.getEntryUsingComparator

源码如下所示：

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

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

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

问题找到了，解决吧！