Limit groupBy in Java 8

Question

How can I limit groupBy by each entry?

For example (based on this example: stream groupBy):

studentClasses.add(new StudentClass("Kumar", 101, "Intro to Web"));
studentClasses.add(new StudentClass("White", 102, "Advanced Java"));
studentClasses.add(new StudentClass("Kumar", 101, "Intro to Cobol"));
studentClasses.add(new StudentClass("White", 101, "Intro to Web"));
studentClasses.add(new StudentClass("White", 102, "Advanced Web"));
studentClasses.add(new StudentClass("Sargent", 106, "Advanced Web"));
studentClasses.add(new StudentClass("Sargent", 103, "Advanced Web"));
studentClasses.add(new StudentClass("Sargent", 104, "Advanced Web"));
studentClasses.add(new StudentClass("Sargent", 105, "Advanced Web"));

This method return simple group:

   Map<String, List<StudentClass>> groupByTeachers = studentClasses
            .stream().collect(
                    Collectors.groupingBy(StudentClass::getTeacher));

What if I want to limit the returned collections? Let's assume I want only the first N classes for every teacher. How can it be done?

Answer 1

It would be possible to introduce a new collector that limits the number of elements in the resulting list.

This collector will retain the head elements of the list (in encounter order). The accumulator and combiner throw away every elements when the limit has been reached during collection. The combiner code is a little tricky but this has the advantage that no additional elements are added only to be thrown away later.

private static <T> Collector<T, ?, List<T>> limitingList(int limit) {
    return Collector.of(
                ArrayList::new, 
                (l, e) -> { if (l.size() < limit) l.add(e); }, 
                (l1, l2) -> {
                    l1.addAll(l2.subList(0, Math.min(l2.size(), Math.max(0, limit - l1.size()))));
                    return l1;
                }
           );
}

And then use it like this:

Map<String, List<StudentClass>> groupByTeachers = 
       studentClasses.stream()
                     .collect(groupingBy(
                          StudentClass::getTeacher,
                          limitingList(2)
                     ));

Answer 2

You could use collectingAndThen to define a finisher operation on the resulting list. This way you can limit, filter, sort, ... the lists:

int limit = 2;

Map<String, List<StudentClass>> groupByTeachers =
    studentClasses.stream()
                  .collect(
                       groupingBy(
                           StudentClass::getTeacher,
                           collectingAndThen(
                               toList(),
                               l -> l.stream().limit(limit).collect(toList()))));