Java: TreeMap is there a nice way to search for duplicate strings as part of a composite key

Question

I have a TreeMap where the key is a composite key based on two fields. I would like to be able to search the TreeMap for matches on just the 2nd key element - there maybe duplicates of this element. To explain what I am trying to do please see the following,

public class CountySchoolsController {

    static TreeMap<StudentKey, Student> schoolsMap = new TreeMap<>();

    public static void main(String args[]) {
        createSchoolsTreeMap();
        System.out.println(schoolsMap.get(new StudentKey(1, "Holmes")));
    }

    private static TreeMap<StudentKey, Student> createSchoolsTreeMap() {
        Student s1 = new Student(1, "Sherlock", "Holmes");
        Student s2 = new Student(2, "John", "Watson");
        Student s3 = new Student(3, "Franklin", "Holmes");
        schoolsMap.put(new StudentKey(s1.getSchoolId(), s1.getLastname()), s1);
        schoolsMap.put(new StudentKey(s2.getSchoolId(), s2.getLastname()), s2);
        schoolsMap.put(new StudentKey(s3.getSchoolId(), s3.getLastname()), s3);
        return schoolsMap;
    }
}

public class StudentKey implements Comparable<StudentKey>{

    int schoolId;
    String lastname;

    public StudentKey(int id, String lastname){
        this.schoolId = id;
        this.lastname = lastname;
    }

    public int getSchoolId() {
        return schoolId;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        StudentKey that = (StudentKey) o;
        return schoolId == that.schoolId &&
                Objects.equals(lastname, that.lastname);
    }

    @Override
    public int hashCode() {
        return Objects.hash(schoolId, lastname);
    }

    @Override
    public int compareTo(StudentKey o) {
        return (this.schoolId + this.lastname).compareTo(o.schoolId + o.lastname);
    }
}

public class Student {

    int schoolId;
    String firstname;
    String lastname;

    public Student(int schoolId, String firstname, String lastname) {
        this.schoolId = schoolId;
        this.firstname = firstname;
        this.lastname = lastname;
    }

    public int getSchoolId() {
        return schoolId;
    }

    public String getFirstname() {
        return firstname;
    }

    public String getLastname() {
        return lastname;
    }

    @Override
    public String toString() {
        return "Student{" +
                "schoolId=" + schoolId +
                ", firstname='" + firstname + '\'' +
                ", lastname='" + lastname + '\'' +
                '}';
    }
}

When I run this code it runs fine and prints the found Student,

Student{schoolId=1, firstname='Sherlock', lastname='Holmes'}

What I would like to be able to do however is search just the lastname for Holmes and return two records represented by Ids 1 and 3. As well as doing searches likes this I also need the ability to do searches on an exact match on the key (as in the example above).

Ideally it would be nice if I could use a wildcard for schoolId but I don't think this is possible.

I could return the keyset values and iterate over this to find the match on just the lastname but I dont think this would be very performant - please tell me if you disagree or if this would be the best way to implement this ? Or should I be implementing this another way ?

Answer 1

Try this. It streams the entrySet of the map, filtering only on the last name, then maps to the values associated with that name and puts it in a list. I had to make the lastname field public for this. Putting in getters for the fields would be useful.

        List<Student> list = schoolsMap.entrySet().stream()
                .filter(e -> e.getKey().lastname
                        .equals("Holmes"))
                .map(Entry::getValue)
                .collect(Collectors.toList());


        list.forEach(System.out::println);

I decided to go a little further on this. If you put getters for all your key attributes in your StudentKey class, you can do the following:

Get all the last names for Holmes

      List<Student> names = getStudentsForKeyAttribute(
                StudentKey::getLastName, "Holmes");

      names.forEach(System.out::println);

Get the student for id = 3

      List<Student> ids = getStudentsForKeyAttribute(StudentKey::getSchoolId, 3);

      ids.forEach(System.out.println); 

The following method accepts an AttributeExtractor function to pull the appropriate attribute from the StudentKey and then filter based on the supplied argument.

        public <T> List<Student> getStudentsForKeyAttribute(
            Function<StudentKey, T> attrExtractor, T keyAttribute) {
            return schoolsMap.entrySet().stream()
                .filter(e -> attrExtractor.apply(e.getKey())
                        .equals(keyAttribute))
                .map(Entry::getValue)
                .collect(Collectors.toList());
        }

MAJOR EDIT: I added data retention capabilities. First time thru for each attribute it builds the map for that attribute and returns requested value. Future calls use existing map.

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;
import java.util.function.Function;


// Main attribute for lookup
enum Attribute {
    LASTNAME, GENDER
}

// subattribute for gender.  For lastname it would jut be a name.
// I decided to use an enum for gender.
enum Gender {
    M, F, O
};

public class CompositeSearch {
    // Map of map.
    // Example:
    //   Outer Map is for LastName attribute
    //   Inner map as all lists for common last names.  All names are included.
    //     I had to use an Object to allow for different types (enums, strings, ints)
    Map<Attribute, Map<Object, List<Student>>> studentsByAttribute = new HashMap<>();


    // this provides an extractor for each type requested. It just maps the 
    // Attribute to the Student Key method call for that type.
    Map<Attribute, Function<StudentKey, ?>> extractorsForType = new HashMap<>() {
        {
            put(Attribute.LASTNAME,
                    StudentKey::getLastName);
            put(Attribute.GENDER, StudentKey::getGender);
        }
    };

    // intiial data base
    TreeMap<StudentKey, Student> schoolsMap = new TreeMap<>();

    public static void main(String args[]) {
        new CompositeSearch().start();
    }

    public void start() {
        createSchoolsTreeMap();

        // getting all female students.
        List<Student> list = getStudentsForKeyAttribute(
                Attribute.GENDER, Gender.F);

        list.forEach(System.out::println);
        System.out.println();

        // getting all students with last name of holmes.
        list = getStudentsForKeyAttribute(Attribute.LASTNAME, "Holmes");
        list.forEach(System.out::println);

        // All maps for Gender and lastnames have been created so 
        // the lookups below require two map retrievals.  The attribute and the 
        // sub attribute
        System.out.println();
        list = getStudentsForKeyAttribute(
                Attribute.GENDER, Gender.M);

        list.forEach(System.out::println);
        System.out.println();

        list = getStudentsForKeyAttribute(Attribute.LASTNAME, "Watson");
        list.forEach(System.out::println);


    }

    public <T> List<Student> getStudentsForKeyAttribute(
            Attribute attr, T keyAttribute) {

        @SuppressWarnings("unchecked")
        Function<StudentKey, T> extractor = (Function<StudentKey, T>) extractorsForType
                .get(attr);

        if (!studentsByAttribute.containsKey(attr)) {
            // need to create the map.
            System.out.println("Building map for all " + attr);
            // sub attribute map
            Map<Object, List<Student>> subMap = new HashMap<>();

            studentsByAttribute.put(attr, subMap);

            for (Map.Entry<StudentKey, ?> e : schoolsMap
                    .entrySet()) {
              T subAttribute = extractor.apply(e.getKey());

                            subMap.compute(subAttribute,
                                    (k, v) -> v == null
                                            ?  new ArrayList<>()
                                            : v)
                            .add((Student)e.getValue());

            }
        } else {
            System.out.println("Using existing map for all " + attr);
        }
        return studentsByAttribute.get(attr).get(keyAttribute);

    }

    // from here on out, everything is pretty normal.
    private TreeMap<StudentKey, Student> createSchoolsTreeMap() {
        List<Student> students = List.of(
                new Student(1, "Sherlock", "Holmes",
                        Gender.M),
                new Student(2, "John", "Watson", Gender.M),
                new Student(3, "Franklin", "Holmes",
                        Gender.M),
                new Student(4, "Frances", "Holmes",
                        Gender.F),
                new Student(5, "Mary", "Wilson", Gender.F),
                new Student(6, "Martha", "Watson",
                        Gender.F));
        for (Student s : students) {
            schoolsMap.put(new StudentKey(s), s);
        }

        return schoolsMap;
    }

}

class StudentKey implements Comparable<StudentKey> {

    private int schoolId;
    private String lastname;
    private Gender gender;

    public StudentKey(Student student) {
        this.schoolId = student.getSchoolId();
        this.lastname = student.getLastname();
        this.gender = student.getGender();
    }

    public int getSchoolId() {
        return schoolId;
    }

    public String getLastName() {
        return lastname;
    }

    public Gender getGender() {
        return gender;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o)
            return true;
        if (o == null || getClass() != o.getClass())
            return false;
        StudentKey that = (StudentKey) o;
        return schoolId == that.schoolId
                && Objects.equals(lastname, that.lastname);
    }

    @Override
    public int hashCode() {
        return Objects.hash(schoolId, lastname);
    }

    @Override
    public int compareTo(StudentKey o) {
        return (this.schoolId + this.lastname)
                .compareTo(o.schoolId + o.lastname);
    }

}

class Student {

    int schoolId;
    String firstname;
    String lastname;
    Gender gender;

    public Student(int schoolId, String firstname,
            String lastname, Gender gender) {
        this.schoolId = schoolId;
        this.firstname = firstname;
        this.lastname = lastname;
        this.gender = gender;
    }

    public int getSchoolId() {
        return schoolId;
    }

    public String getFirstname() {
        return firstname;
    }

    public String getLastname() {
        return lastname;
    }

    public Gender getGender() {
        return gender;
    }

    @Override
    public String toString() {
        return "Student{" + "schoolId=" + schoolId
                + ", firstname='" + firstname + '\''
                + ", lastname='" + lastname + '\'' + '}';
    }
}

Answer 2

Well, filtering is cool, but slow. So if you want some performance - let's add index:

Map<String,Set<StudentKey>> lastNamesMap = new HashMap<>();

Set<StudentKey> getByLastName(String lastName) 
  return lastNamesMap.containsKey(s1.getLastname()) ? lastNamesMap.get(s1.getLastname()) : Collections.emptySet();  
}

void addStudent(Student s1) {
 final String k = s1.getLastname();
 Set<StudentKey> keys;
 if (lastNamesMap.containsKey(k)) {
  keys = lastNamesMap.get(k);
 } else {
  keys = new TreeSet<>();
 }
 StudentKey key = new StudentKey(s1.getSchoolId(), s1.getLastname();
 keys.add(key);
 lastNamesMap.put(k,keys);
 schoolsMap.put(key, s1);
 }