Flatten a map after Collectors.groupingBy in java

Question

I have list of students. I want to return list of objects StudentResponse classes that has the course and the list of students for the course. So I can write which gives me a map

Map<String, List<Student>> studentsMap = students.stream().
            .collect(Collectors.groupingBy(Student::getCourse,
                    Collectors.mapping(s -> s, Collectors.toList()
             )));

Now I have to iterate through the map again to create a list of objects of StudentResponse class which has the Course and List:

class StudentResponse {
     String course;
     Student student;

     // getter and setter
}

Is there a way to combine these two iterations?

Answer 1

Not exactly what you've asked, but here's a compact way to accomplish what you want, just for completeness:

Map<String, StudentResponse> map = new LinkedHashMap<>();
students.forEach(s -> map.computeIfAbsent(
        s.getCourse(), 
        k -> new StudentResponse(s.getCourse()))
    .getStudents().add(s));

This assumes StudentResponse has a constructor that accepts the course as an argument and a getter for the student list, and that this list is mutable (i.e. ArrayList) so that we can add the current student to it.

While the above approach works, it clearly violates a fundamental OO principle, which is encapsulation. If you are OK with that, then you're done. If you want to honor encapsulation, then you could add a method to StudentResponse to add a Student instance:

public void addStudent(Student s) {
    students.add(s);
}

Then, the solution would become:

Map<String, StudentResponse> map = new LinkedHashMap<>();
students.forEach(s -> map.computeIfAbsent(
        s.getCourse(), 
        k -> new StudentResponse(s.getCourse()))
    .addStudent(s));

This solution is clearly better than the previous one and would avoid a rejection from a serious code reviewer.

Both solutions rely on Map.computeIfAbsent, which either returns a StudentResponse for the provided course (if there exists an entry for that course in the map), or creates and returns a StudentResponse instance built with the course as an argument. Then, the student is being added to the internal list of students of the returned StudentResponse.

Finally, your StudentResponse instances are in the map values:

Collection<StudentResponse> result = map.values();

If you need a List instead of a Collection:

List<StudentResponse> result = new ArrayList<>(map.values());

Note: I'm using LinkedHashMap instead of HashMap to preserve insertion-order, i.e. the order of the students in the original list. If you don't have such requirement, just use HashMap.

Answer 2

Probably way overkill but it was a fun exercise :) You could implement your own Collector:

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.*;
import java.util.stream.Collector;
import java.util.stream.Collectors;

public class StudentResponseCollector implements Collector<Student, Map<String, List<Student>>, List<StudentResponse>> {

    @Override
    public Supplier<Map<String, List<Student>>> supplier() {
        return () -> new ConcurrentHashMap<>();
    }

    @Override
    public BiConsumer<Map<String, List<Student>>, Student> accumulator() {
        return (store, student) -> store.merge(student.getCourse(),
                new ArrayList<>(Arrays.asList(student)), combineLists());
    }

    @Override
    public BinaryOperator<Map<String, List<Student>>> combiner() {
        return (x, y) -> {
            x.forEach((k, v) -> y.merge(k, v, combineLists()));

            return y;
        };
    }

    private <T> BiFunction<List<T>, List<T>, List<T>> combineLists() {
        return (students, students2) -> {
            students2.addAll(students);
            return students2;
        };
    }

    @Override
    public Function<Map<String, List<Student>>, List<StudentResponse>> finisher() {
        return (store) -> store
                .keySet()
                .stream()
                .map(course -> new StudentResponse(course, store.get(course)))
                .collect(Collectors.toList());
    }

    @Override
    public Set<Characteristics> characteristics() {
        return EnumSet.of(Characteristics.UNORDERED);
    }
}

Given Student and StudentResponse:

public class Student {
    private String name;
    private String course;

    public Student(String name, String course) {
        this.name = name;
        this.course = course;
    }

    public String getName() {
        return name;
    }

    public String getCourse() {
        return course;
    }

    public String toString() {
        return name + ", " + course;
    }
}

public class StudentResponse {
    private String course;
    private List<Student> studentList;

    public StudentResponse(String course, List<Student> studentList) {
        this.course = course;
        this.studentList = studentList;
    }

    public String getCourse() {
        return course;
    }

    public List<Student> getStudentList() {
        return studentList;
    }

    public String toString() {
        return course + ", " + studentList.toString();
    }
}

Your code where you collect your StudentResponses can now be very short and elegant ;)

public class StudentResponseCollectorTest {

    @Test
    public void test() {
        Student student1 = new Student("Student1", "foo");
        Student student2 = new Student("Student2", "foo");
        Student student3 = new Student("Student3", "bar");

        List<Student> studentList = Arrays.asList(student1, student2, student3);

        List<StudentResponse> studentResponseList = studentList
                .stream()
                .collect(new StudentResponseCollector());

        assertEquals(2, studentResponseList.size());
    }
}