Java TreeMap e TreeSet

O HashMap e o HashSet são rápidos, mas embaralham os elementos. Quando você quer as coisas sempre em ordem (os alunos de A a Z, as pontuações da menor para a maior) e ainda perguntar "qual é a primeira?", "quem vem logo depois da Bia?", use o TreeMap e o TreeSet. Eles se reorganizam sozinhos a cada elemento que entra.

A ideia: um fichário em ordem alfabética

A secretaria da escola guarda uma ficha por aluno num fichário. Toda ficha nova é encaixada no lugar certo, em ordem alfabética. Assim, ninguém precisa reordenar nada depois, e é fácil responder "qual é a primeira ficha?" ou "qual vem logo depois de 'Bia'?".

Por dentro, essas coleções usam uma estrutura chamada árvore (tree, em inglês; mais precisamente, uma "árvore rubro-negra"), que mantém os elementos ordenados e permite achar qualquer um com poucos passos. Daí o nome: Tree + Map, Tree + Set. Elas são um pouco mais lentas que as versões Hash, mas a diferença só aparece com quantidades muito grandes.

TreeMap: um mapa com chaves ordenadas

O TreeMap é um Map como qualquer outro: put, get, remove, containsKey… tudo igual ao HashMap. A diferença é a ordem e os métodos extras de navegação:

import java.util.TreeMap;

public class Main {
    public static void main(String[] args) {
        TreeMap<String, Double> notas = new TreeMap<>();
        notas.put("Carla", 8.25);
        notas.put("Ana", 9.5);
        notas.put("Edu", 6.0);
        notas.put("Bruno", 7.0);
        notas.put("Duda", 10.0);
        System.out.println(notas);

        System.out.println("firstKey(): " + notas.firstKey());
        System.out.println("lastKey():  " + notas.lastKey());
        System.out.println("firstEntry(): " + notas.firstEntry());

        System.out.println("headMap(\"Carla\"): " + notas.headMap("Carla"));
        System.out.println("tailMap(\"Carla\"): " + notas.tailMap("Carla"));
        System.out.println("subMap(\"B\", \"D\"): " + notas.subMap("B", "D"));
    }
}

$ javac Main.java
$ java Main
{Ana=9.5, Bruno=7.0, Carla=8.25, Duda=10.0, Edu=6.0}      « em ordem, embora inseridos fora de ordem
firstKey(): Ana
lastKey():  Edu
firstEntry(): Ana=9.5
headMap("Carla"): {Ana=9.5, Bruno=7.0}                     « antes de Carla (sem ela)
tailMap("Carla"): {Carla=8.25, Duda=10.0, Edu=6.0}         « de Carla em diante (com ela)
subMap("B", "D"): {Bruno=7.0, Carla=8.25}                  « de "B" (incluso) até "D" (excluso)

TreeMap<String, Double> notas

Aqui a variável é TreeMap, e não Map, porque os métodos de navegação (firstKey, headMap…) só existem nos mapas ordenados. Se você só precisa da ordem ao imprimir, Map<…> notas = new TreeMap<>() basta.

firstKey() / lastKey()

A menor / a maior chave.

firstEntry() / lastEntry()

O par inteiro (chave=valor) da ponta.

headMap(k)

"A cabeça": os pares com chave menor que k. O k fica de fora.

tailMap(k)

"A cauda": os pares com chave maior ou igual a k.

subMap(de, ate)

Uma "fatia": de de (incluso) até ate (excluso). subMap("B", "D") pega todo nome que começa com B ou C.

A regra "começo incluso, fim excluso" é a mesma do substring (veja Métodos de String). Ela tem uma vantagem: headMap("Carla") e tailMap("Carla") dividem o mapa em duas partes sem repetir ninguém.

Buscar o vizinho: ceiling, floor, higher, lower

Às vezes a chave que você procura não está no mapa, e você quer a mais próxima. Pense num elevador: ceiling = teto (para cima), floor = chão (para baixo).

        System.out.println("ceilingKey(\"Bia\"):  " + notas.ceilingKey("Bia"));
        System.out.println("ceilingKey(\"Carla\"): " + notas.ceilingKey("Carla"));
        System.out.println("higherKey(\"Carla\"):  " + notas.higherKey("Carla"));
        System.out.println("floorKey(\"Bia\"):    " + notas.floorKey("Bia"));
        System.out.println("lowerKey(\"Ana\"):    " + notas.lowerKey("Ana"));
        System.out.println("ceilingKey(\"Zeca\"): " + notas.ceilingKey("Zeca"));

Com o mapa {Ana, Bruno, Carla, Edu}:

$ java Main
ceilingKey("Bia"):  Bruno           « a menor chave >= "Bia"
ceilingKey("Carla"): Carla          « Carla existe: ela mesma
higherKey("Carla"):  Edu            « estritamente maior: pula a própria Carla
floorKey("Bia"):    Ana             « a maior chave <= "Bia"
lowerKey("Ana"):    null            « não há ninguém antes da Ana
ceilingKey("Zeca"): null            « não há ninguém depois do Zeca

Todos devolvem null quando não existe vizinho naquela direção. E cada um tem uma versão …Entry (ceilingEntry, floorEntry…) que devolve o par inteiro.

Um uso real: faixas de nota

A escola da Ana converte notas em conceitos: de 9 para cima é A, de 7 a 9 é B, de 5 a 7 é C, abaixo de 5 é E. Com um TreeMap cujas chaves são o início de cada faixa, o floorEntry acha a faixa de qualquer nota, sem nenhum if:

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

public class Main {
    public static void main(String[] args) {
        TreeMap<Double, String> conceitos = new TreeMap<>();
        conceitos.put(0.0, "E");
        conceitos.put(5.0, "C");
        conceitos.put(7.0, "B");
        conceitos.put(9.0, "A");

        double[] notas = {9.5, 6.5, 7.0, 3.0};
        for (double n : notas) {
            Map.Entry<Double, String> faixa = conceitos.floorEntry(n);
            System.out.println(n + " -> faixa a partir de " + faixa.getKey() + " = " + faixa.getValue());
        }
    }
}

$ javac Main.java
$ java Main
9.5 -> faixa a partir de 9.0 = A
6.5 -> faixa a partir de 5.0 = C       « o "chão" de 6.5 é 5.0
7.0 -> faixa a partir de 7.0 = B       « 7.0 existe: é ele mesmo
3.0 -> faixa a partir de 0.0 = E

Para mudar as faixas, você muda os dados do mapa, não o código. A mesma ideia serve para faixas de imposto, de frete por peso, de desconto por quantidade…

import java.util.TreeMap;

public class Main {
    public static void main(String[] args) {
        TreeMap<Double, String> conceitos = new TreeMap<>();
        conceitos.put(0.0, "E");
        conceitos.put(5.0, "C");
        conceitos.put(7.0, "B");
        conceitos.put(9.0, "A");

        double nota = 8.0;
        System.out.println("Conceito: " + conceitos.get(nota));
    }
}

import java.util.TreeMap;

public class Main {
    public static void main(String[] args) {
        TreeMap<Double, String> conceitos = new TreeMap<>();
        conceitos.put(0.0, "E");
        conceitos.put(5.0, "C");
        conceitos.put(7.0, "B");
        conceitos.put(9.0, "A");

        double nota = 8.0;
        System.out.println("Conceito: " + conceitos.floorEntry(nota).getValue());
    }
}

TreeSet: um conjunto sempre ordenado

O TreeSet é para o TreeMap o que o HashSet é para o HashMap: só as chaves. É um conjunto (sem repetidos) que mantém os elementos em ordem, com os mesmos métodos de navegação, sem o "Key" no nome:

import java.util.TreeSet;

public class Main {
    public static void main(String[] args) {
        TreeSet<Integer> pontos = new TreeSet<>();
        pontos.add(70);
        pontos.add(15);
        pontos.add(42);
        pontos.add(99);
        pontos.add(42);
        pontos.add(8);
        System.out.println(pontos);

        System.out.println("first(): " + pontos.first() + "  last(): " + pontos.last());
        System.out.println("headSet(42): " + pontos.headSet(42));
        System.out.println("tailSet(42): " + pontos.tailSet(42));
        System.out.println("ceiling(50): " + pontos.ceiling(50) + "  floor(50): " + pontos.floor(50));
        System.out.println("higher(99): " + pontos.higher(99));
        System.out.println("descendingSet(): " + pontos.descendingSet());
        System.out.println("pollFirst(): " + pontos.pollFirst() + " -> " + pontos);
    }
}

$ javac Main.java
$ java Main
[8, 15, 42, 70, 99]                       « ordenado e sem o 42 repetido
first(): 8  last(): 99
headSet(42): [8, 15]
tailSet(42): [42, 70, 99]
ceiling(50): 70  floor(50): 42
higher(99): null
descendingSet(): [99, 70, 42, 15, 8]      « a mesma coisa, de trás para a frente
pollFirst(): 8 -> [15, 42, 70, 99]        « tira e devolve o menor

descendingSet() não cria uma cópia: é uma "visão" ao contrário do mesmo conjunto. pollFirst()/pollLast() tiram o menor/o maior, o que transforma o TreeSet numa fila de prioridade simples.

Ordem natural × Comparator

Como o TreeSet sabe que 8 vem antes de 15 e "Ana" antes de "Bruno"? Pela ordem natural do tipo: números do menor para o maior; textos pelos códigos dos caracteres (maiúsculas antes das minúsculas; veja a aula de ArrayList). Os tipos que têm ordem natural implementam a interface Comparable.

Quando a ordem natural não serve, você passa ao construtor um Comparator (comparador): um objeto que diz como comparar dois elementos. Veja as mesmas frutas com quatro critérios:

import java.util.Comparator;
import java.util.TreeSet;

public class Main {
    public static void main(String[] args) {
        String[] frutas = {"uva", "Banana", "kiwi", "abacaxi", "pera", "Caju"};

        TreeSet<String> natural = new TreeSet<>();
        TreeSet<String> semCaixa = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
        TreeSet<String> inversa = new TreeSet<>(Comparator.reverseOrder());
        TreeSet<String> porTamanho = new TreeSet<>(Comparator.comparing(String::length));
        for (String f : frutas) {
            natural.add(f);
            semCaixa.add(f);
            inversa.add(f);
            porTamanho.add(f);
        }
        System.out.println("natural:     " + natural);
        System.out.println("sem caixa:   " + semCaixa);
        System.out.println("inversa:     " + inversa);
        System.out.println("por tamanho: " + porTamanho);
    }
}

new TreeSet<>()

Sem argumento: ordem natural.

String.CASE_INSENSITIVE_ORDER

Um comparador pronto que ignora maiúsculas/minúsculas.

Comparator.reverseOrder()

A ordem natural ao contrário.

Comparator.comparing(String::length)

"Compare pelo tamanho" (length). O :: é uma referência de método.

$ javac Main.java
$ java Main
natural:     [Banana, Caju, abacaxi, kiwi, pera, uva]      « maiúsculas primeiro
sem caixa:   [abacaxi, Banana, Caju, kiwi, pera, uva]      « como num dicionário
inversa:     [uva, pera, kiwi, abacaxi, Caju, Banana]
por tamanho: [uva, kiwi, Banana, abacaxi]                  « sumiram "pera" e "Caju"!

Quando usar cada um

Quando dá errado

import java.util.TreeSet;

class Aluno {
    String nome;
    Aluno(String nome) { this.nome = nome; }
}

public class Main {
    public static void main(String[] args) {
        TreeSet<Aluno> turma = new TreeSet<>();
        turma.add(new Aluno("Ana"));
    }
}

$ javac Main.java
$ java Main
Exception in thread "main" java.lang.ClassCastException: class Aluno cannot be cast to class java.lang.Comparable (Aluno 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:1605)
	at java.base/java.util.TreeMap.addEntryToEmptyMap(TreeMap.java:810)
	at java.base/java.util.TreeMap.put(TreeMap.java:819)
	at java.base/java.util.TreeMap.put(TreeMap.java:568)
	at java.base/java.util.TreeSet.add(TreeSet.java:261)
	at Main.main(Main.java:11)

$ java Main
Exception in thread "main" java.lang.NullPointerException
	at java.base/java.util.Objects.requireNonNull(Objects.java:220)
	at java.base/java.util.TreeMap.put(TreeMap.java:843)
	at java.base/java.util.TreeMap.put(TreeMap.java:568)
	at Main.main(Main.java:7)

$ java Main
Exception in thread "main" java.util.NoSuchElementException
	at java.base/java.util.TreeMap.key(TreeMap.java:1638)
	at java.base/java.util.TreeMap.firstKey(TreeMap.java:303)
	at Main.main(Main.java:6)

Referência

Exercícios

import java.util.TreeSet;

public class Main {
    public static void main(String[] args) {
        TreeSet<String> frutas = new TreeSet<>();
        frutas.add("banana");
        frutas.add("Abacaxi");
        frutas.add("uva");
        frutas.add("Caju");
        System.out.println(frutas);
        System.out.println("Primeira: " + frutas.first());
    }
}

import java.util.TreeSet;

public class Main {
    public static void main(String[] args) {
        TreeSet<String> frutas = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
        frutas.add("banana");
        frutas.add("Abacaxi");
        frutas.add("uva");
        frutas.add("Caju");
        System.out.println(frutas);
        System.out.println("Primeira: " + frutas.first());
    }
}

Resumo

• TreeMap = mapa com chaves sempre ordenadas; TreeSet = conjunto sempre ordenado.
• Pontas: firstKey/lastKey (mapa), first/last (conjunto).
• Fatias: headMap/headSet (antes, sem incluir), tailMap/tailSet (dali em diante), subMap/subSet.
• Vizinhos: ceiling (>=), higher (>), floor (<=), lower (<); null se não existir. floorEntry resolve faixas.
• Sem argumento, usam a ordem natural (Comparable); com um Comparator, a ordem que você quiser.
• Elementos que o comparador acha "iguais" viram um só. Não aceitam null. Classe sem ordem natural: ClassCastException.

Próxima aula: Iterator, o "cursor" que percorre qualquer coleção, e como remover elementos durante o percurso.