MK
摩柯社区 - 一个极简的技术知识社区
AI 面试

如何在Java TreeMap中自定义比较器

2023-12-315.2k 阅读

Java TreeMap基础回顾

在深入探讨如何在Java TreeMap中自定义比较器之前，我们先来回顾一下TreeMap的基础知识。TreeMap是Java集合框架中的一部分，它实现了NavigableMap接口，基于红黑树（Red - Black Tree）数据结构来存储键值对。这使得TreeMap中的键是有序的，默认情况下，键按照自然顺序（natural order）进行排序。

例如，如果键是Integer类型，它们会按照从小到大的顺序排列；如果键是String类型，则按照字典序排列。

下面是一个简单的使用TreeMap默认排序的示例代码：

import java.util.Map;
import java.util.TreeMap;

public class TreeMapDefaultOrderExample {
    public static void main(String[] args) {
        Map<Integer, String> treeMap = new TreeMap<>();
        treeMap.put(3, "Three");
        treeMap.put(1, "One");
        treeMap.put(2, "Two");

        for (Map.Entry<Integer, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey() + " : " + entry.getValue());
        }
    }
}

在上述代码中，我们创建了一个TreeMap，并向其中添加了一些键值对。当我们遍历这个TreeMap时，会发现输出的键是按照从小到大的顺序排列的，即：

1 : One
2 : Two
3 : Three

为什么需要自定义比较器

虽然TreeMap的默认自然排序在很多情况下都能满足需求，但在实际应用中，我们常常会遇到一些复杂的场景，自然排序无法满足要求。

比如，我们有一个自定义的类作为TreeMap的键，这个类可能没有实现Comparable接口，或者我们希望按照与自然排序不同的规则来对键进行排序。

假设有一个Person类，包含name和age两个属性，我们可能希望根据age来对Person对象进行排序，而不是使用默认的对象引用比较（因为Person类没有实现Comparable接口时，默认是比较对象的内存地址）。

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

如果我们尝试直接将Person对象作为键放入TreeMap中，会得到一个ClassCastException，因为Person类没有实现Comparable接口，TreeMap不知道如何比较两个Person对象。

import java.util.Map;
import java.util.TreeMap;

public class PersonAsKeyWithoutComparatorExample {
    public static void main(String[] args) {
        Map<Person, String> treeMap = new TreeMap<>();
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");
    }
}

运行上述代码会抛出如下异常：

Exception in thread "main" java.lang.ClassCastException: class Person cannot be cast to class java.lang.Comparable (Person is in unnamed module of loader 'app'; java.lang.Comparable is in module java.base of loader 'bootstrap')
	at java.base/java.util.TreeMap.compare(TreeMap.java:1294)
	at java.base/java.util.TreeMap.put(TreeMap.java:538)
	at PersonAsKeyWithoutComparatorExample.main(PersonAsKeyWithoutComparatorExample.java:8)

这时候，我们就需要自定义比较器来告诉TreeMap如何比较两个Person对象。

使用匿名内部类自定义比较器

一种常见的自定义比较器的方式是使用匿名内部类。Java提供了Comparator接口，我们可以实现这个接口来自定义比较逻辑。

继续以上面的Person类为例，我们可以使用匿名内部类来创建一个比较器，按照age对Person对象进行升序排序。

import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

public class TreeMapCustomComparatorAnonymousInnerClassExample {
    public static void main(String[] args) {
        Comparator<Person> ageComparator = new Comparator<Person>() {
            @Override
            public int compare(Person p1, Person p2) {
                return p1.getAge() - p2.getAge();
            }
        };

        Map<Person, String> treeMap = new TreeMap<>(ageComparator);
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");

        for (Map.Entry<Person, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey().getName() + " : " + entry.getKey().getAge() + " : " + entry.getValue());
        }
    }
}

在上述代码中，我们首先创建了一个实现了Comparator接口的匿名内部类ageComparator。在compare方法中，我们通过比较两个Person对象的age属性来确定它们的顺序。如果p1的age小于p2的age，返回一个负整数；如果p1的age大于p2的age，返回一个正整数；如果age相等，返回0。

然后，我们将这个比较器传递给TreeMap的构造函数，这样TreeMap就会按照我们定义的比较器来对Person对象进行排序。运行结果如下：

Bob : 20 : Info about Bob
Alice : 25 : Info about Alice

使用Lambda表达式自定义比较器

从Java 8开始，Lambda表达式为我们提供了一种更简洁的方式来实现Comparator接口。我们可以将上述使用匿名内部类的代码改写成使用Lambda表达式的形式。

import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

public class TreeMapCustomComparatorLambdaExample {
    public static void main(String[] args) {
        Comparator<Person> ageComparator = (p1, p2) -> p1.getAge() - p2.getAge();

        Map<Person, String> treeMap = new TreeMap<>(ageComparator);
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");

        for (Map.Entry<Person, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey().getName() + " : " + entry.getKey().getAge() + " : " + entry.getValue());
        }
    }
}

在上述代码中，Comparator<Person> ageComparator = (p1, p2) -> p1.getAge() - p2.getAge();这一行代码使用Lambda表达式创建了一个比较器。Lambda表达式(p1, p2) -> p1.getAge() - p2.getAge()等价于匿名内部类中的compare方法的实现。这种方式更加简洁明了，代码量更少。运行结果与使用匿名内部类时相同：

Bob : 20 : Info about Bob
Alice : 25 : Info about Alice

自定义复杂比较逻辑

有时候，我们的比较逻辑可能会更加复杂。比如，对于Person类，我们可能希望首先按照age进行排序，如果age相同，再按照name进行字典序排序。

我们可以通过以下方式实现：

import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

public class TreeMapComplexComparatorExample {
    public static void main(String[] args) {
        Comparator<Person> complexComparator = (p1, p2) -> {
            int ageComparison = p1.getAge() - p2.getAge();
            if (ageComparison != 0) {
                return ageComparison;
            } else {
                return p1.getName().compareTo(p2.getName());
            }
        };

        Map<Person, String> treeMap = new TreeMap<>(complexComparator);
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        Person person3 = new Person("Charlie", 25);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");
        treeMap.put(person3, "Info about Charlie");

        for (Map.Entry<Person, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey().getName() + " : " + entry.getKey().getAge() + " : " + entry.getValue());
        }
    }
}

在上述代码中，complexComparator首先比较两个Person对象的age。如果age不同，直接返回age的比较结果。如果age相同，再通过name的compareTo方法进行字典序比较。运行结果如下：

Bob : 20 : Info about Bob
Alice : 25 : Info about Alice
Charlie : 25 : Info about Charlie

基于已有比较器构建新比较器

Java的Comparator接口提供了一些静态方法，我们可以基于已有比较器构建新的比较器，这进一步简化了复杂比较逻辑的实现。

例如，我们可以使用Comparator.comparing方法来创建一个基于对象某个属性的比较器，然后使用thenComparing方法来添加额外的比较逻辑。

import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

public class TreeMapChainedComparatorExample {
    public static void main(String[] args) {
        Comparator<Person> chainedComparator = Comparator.comparing(Person::getAge)
               .thenComparing(Person::getName);

        Map<Person, String> treeMap = new TreeMap<>(chainedComparator);
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        Person person3 = new Person("Charlie", 25);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");
        treeMap.put(person3, "Info about Charlie");

        for (Map.Entry<Person, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey().getName() + " : " + entry.getKey().getAge() + " : " + entry.getValue());
        }
    }
}

在上述代码中，Comparator.comparing(Person::getAge)首先创建了一个基于age属性的比较器。然后，通过thenComparing(Person::getName)添加了如果age相同则按照name进行字典序比较的逻辑。运行结果与前面复杂比较逻辑的示例相同：

Bob : 20 : Info about Bob
Alice : 25 : Info about Alice
Charlie : 25 : Info about Charlie

反向排序

有时候，我们可能希望按照与默认顺序相反的顺序进行排序。对于简单类型，我们可以使用Collections.reverseOrder方法来获取反向比较器。对于自定义类型，我们可以在比较器的compare方法中反转比较逻辑。

以Person类按照age降序排序为例，使用Lambda表达式可以这样实现：

import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public String getName() {
        return name;
    }

    public int getAge() {
        return age;
    }
}

public class TreeMapReverseOrderExample {
    public static void main(String[] args) {
        Comparator<Person> reverseAgeComparator = (p1, p2) -> p2.getAge() - p1.getAge();

        Map<Person, String> treeMap = new TreeMap<>(reverseAgeComparator);
        Person person1 = new Person("Alice", 25);
        Person person2 = new Person("Bob", 20);
        treeMap.put(person1, "Info about Alice");
        treeMap.put(person2, "Info about Bob");

        for (Map.Entry<Person, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey().getName() + " : " + entry.getKey().getAge() + " : " + entry.getValue());
        }
    }
}

运行结果如下：

Alice : 25 : Info about Alice
Bob : 20 : Info about Bob

如果我们想对简单类型（如Integer）进行反向排序，可以使用Collections.reverseOrder方法：

import java.util.Collections;
import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;

public class TreeMapReverseOrderIntegerExample {
    public static void main(String[] args) {
        Comparator<Integer> reverseComparator = Collections.reverseOrder();
        Map<Integer, String> treeMap = new TreeMap<>(reverseComparator);
        treeMap.put(3, "Three");
        treeMap.put(1, "One");
        treeMap.put(2, "Two");

        for (Map.Entry<Integer, String> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey() + " : " + entry.getValue());
        }
    }
}

运行结果如下：

3 : Three
2 : Two
1 : One

自定义比较器与性能

在使用自定义比较器时，需要注意性能问题。复杂的比较逻辑可能会导致比较操作的时间复杂度增加，从而影响TreeMap的插入、查找和删除等操作的性能。

例如，如果比较器中包含大量的计算或者数据库查询等耗时操作，会严重影响TreeMap的性能。因此，在设计比较器时，应尽量保持比较逻辑的简洁高效。

另外，由于TreeMap基于红黑树数据结构，其插入、查找和删除操作的时间复杂度在平均情况下为O(log n)，其中n是TreeMap中元素的个数。但是，如果比较器的实现导致元素的比较次数过多，这个平均时间复杂度可能会受到影响。

总结自定义比较器的要点

  1. 实现Comparator接口：无论是使用匿名内部类还是Lambda表达式，都要实现Comparator接口的compare方法，明确比较逻辑。
  2. 比较逻辑的准确性：确保比较逻辑符合实际需求，对于复杂比较逻辑要仔细测试，避免出现逻辑错误。
  3. 性能考虑：尽量保持比较逻辑的简洁高效，避免在比较器中进行耗时操作。
  4. 利用已有工具：充分利用Java提供的Comparator接口的静态方法，如comparing和thenComparing等，来简化复杂比较逻辑的实现。

通过合理地使用自定义比较器，我们可以让TreeMap按照我们期望的方式对键进行排序，从而满足各种复杂的业务需求。无论是简单的升序、降序排序，还是复杂的多条件排序，都可以通过自定义比较器来实现。在实际开发中，根据具体的业务场景，灵活选择合适的方式来定义比较器，能够提高代码的可读性和可维护性，同时也能保证程序的性能。