java TreeMap 数据丢失

在项目中需要对一个HashMap userMap进行排序,排序规则为：HashSet keySet中存在的key优先级高于其它的key，使用TreeMap进行保存 ，
JDK8 + Lombok1.18

一、代码

public class MapTest {

    @Test
    public void testTreeMap() {
        // 里面保存的数据优先级要高些
        Set keySet = new HashSet<>();
        keySet.add("D") ;
        keySet.add("B") ;

        // 要排序的数据
        Map userMap = new HashMap<>();
        userMap.put("A",new User("A",18));
        userMap.put("B",new User("B",15));
        userMap.put("C",new User("C",18));
        userMap.put("D",new User("D",12));

        // 排序后的map
        Map sortMap = new TreeMap<>((k1, k2) -> {
            boolean exist1 = keySet.contains(k1);
            boolean exist2 = keySet.contains(k2);

           if (!exist1 && exist2) {
                return -1;
            } else if (exist1 && !exist2) {
                return -1;
            }
            return 0;
        });

        sortMap.putAll(userMap);

        System.out.println("");
    }

    @Setter
    @Getter
    @NoArgsConstructor
    @AllArgsConstructor
    public static class User{
        private String name ;
        private Integer age ;
    }
}

以上排序后的结果可能是什么？期望结果是B和D排在前面，其它的排后面

二、非预期输出：

结果如下：

image.png

• 数据丢失了两个。并且数据不对

  
  
  

  image.png

不得不吐槽这个代码块的配色，黑挫挫的，一点都不利于阅读，注释什么的看都看不清！！！！！

三、通过源码查找原因

• 直接打开TreeMap源码部分, 看 put()方法就可以了

    public V put(K key, V value) {
        Entry t = root;
        if (t == null) {
            compare(key, key); // type (and possibly null) check

            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        int cmp;
        Entry parent;
        // split comparator and comparable paths
        Comparator cpr = comparator;
        if (cpr != null) {
            do {
                parent = t;
                cmp = cpr.compare(key, t.key);
                if (cmp < 0)
                    t = t.left;
                else if (cmp > 0)
                    t = t.right;
                else
                    return t.setValue(value);
            } while (t != null);
        }
        else {
            if (key == null)
                throw new NullPointerException();
            @SuppressWarnings("unchecked")
                Comparable k = (Comparable) key;
            do {
                parent = t;
                cmp = k.compareTo(t.key);
                if (cmp < 0)
                    t = t.left;
                else if (cmp > 0)
                    t = t.right;
                else
                    return t.setValue(value);
            } while (t != null);
        }
        Entry e = new Entry<>(key, value, parent);
        if (cmp < 0) 
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
    }

• 因为是自定义了comparator , 所以只看下面部分代码片段

        int cmp;
        Entry parent;
        // split comparator and comparable paths
        Comparator cpr = comparator;
        if (cpr != null) {
            do {
                parent = t;
                cmp = cpr.compare(key, t.key);
                if (cmp < 0)
                    t = t.left;  // 标记1
                else if (cmp > 0)
                    t = t.right;  // 标记2
                else
                    return t.setValue(value); // 标记3
            } while (t != null);
        }

• TreeMap 使用了红黑树结构保存数据。喜欢磕一下红黑树细节的先看这个 红黑树实现分析 。
  标记1 ： comparator.compare(k1, k2) 返回值小于0 ，则把值放在左子树，相当于排在前面.
  标记2 ： comparator.compare(k1, k2) 返回值大于0 ，则把值放在右子树，相当于排在后面.
  标记3 ： comparator.compare(k1, k2) 返回值等于0 ，把k1和k2当成相同key，并替换值。这就造成了上图中key和具体数据不匹配的情况 。

四、解决方法

• 修复BUG : TreeMap 的比较器不返回0

       Map sortMap = new TreeMap<>((k1, k2) -> {
            boolean exist1 = keySet.contains(k1);
            boolean exist2 = keySet.contains(k2);

            if (exist1 && !exist2) {
                return -1;
            }
            return 1;
        });

• 修改后的结果如下：

  
  
  

  image.png
  
  
  

  image.png

五、扩展compareTo

public class MapTest {

    @Test
    public void testTreeMap() {
        Map userMap = new HashMap<>();
        User a = new User("A", 18);
        User b = new User("B", 15);
        User c = new User("C", 18);
        User d = new User("D", 12);
        userMap.put(a, a);
        userMap.put(b, b);
        userMap.put(c, c);
        userMap.put(d, d);

        Map sortMap = new TreeMap<>();

        sortMap.putAll(userMap);

        System.out.println("");
    }

    @Setter
    @Getter
    @NoArgsConstructor
    @AllArgsConstructor
    private static class User implements Comparable{
        private String name ;
        private Integer age ;

        @Override
        public int compareTo(User u) {
            return this.getAge() - u.age ;
        }
    }
}

• 上面这个例子中TreeMap的key使用的是User对象并实现了Comparable接口，通过age的排序。实际排序结果还是会丢数据、数据错乱，解决方案就是compareTo中不要返回0 ，前提条件是user.compareTo()只用于TreeMap排序。

例子中只是为了说明compareTo用于TreeMap排序时也不能返回0，所以就不管Map结构是否合理、hashCode和equals是否规范

---

总结：TreeMap 自定义比较器，Comparator不能返回0 ！！！compareTo()也不能返回0 ！！！仅限于要比较的对象只用于TreeMap才要注意

• 红黑树实现分析