条件ごとにStreamを分類する

StreamUtils.java

package com.example;

import java.util.List;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collector.Characteristics;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static java.util.Arrays.asList;
import static java.util.stream.Collectors.partitioningBy;

public class StreamUtils {
    // https://stackoverflow.com/a/35407379
    public static <T, A, R> Collector<T, ?, List<R>> classifyingBy(
            Collector<? super T, A, R> downstream, Predicate<T>... predicates) {
        Supplier<A> dsSupplier = downstream.supplier();
        BiConsumer<A, ? super T> dsAccumulator = downstream.accumulator();
        BinaryOperator<A> dsCombiner = downstream.combiner();

        // downstreamのCollectorにちょっかいをかける
        // Supplier<A>をpredicates分Listを持つSupplier<List<A>>に変換
        Supplier<List<A>> supplier = () -> Stream.generate(dsSupplier)
                .limit(predicates.length).collect(Collectors.toList());
        // predicatesごとにList<A>にTを分類
        BiConsumer<List<A>, T> accumulator = (list, t) -> IntStream.range(0, predicates.length)
                .filter(i -> predicates[i].test(t))
                .forEach(i -> dsAccumulator.accept(list.get(i), t));
        // predicatesごとにList<A>をマージ
        BinaryOperator<List<A>> combiner = (l1, l2) -> IntStream.range(0, predicates.length)
                .mapToObj(i -> dsCombiner.apply(l1.get(i), l2.get(i)))
                .collect(Collectors.toList());

        Characteristics[] dsCharacteristics = downstream.characteristics().toArray(new Characteristics[0]);
        if (downstream.characteristics().contains(Characteristics.IDENTITY_FINISH)) {
            // Finisherが省略可能の場合
            @SuppressWarnings("unchecked")
            Collector<T, ?, List<R>> result = (Collector<T, ?, List<R>>) (Collector<T, ?, ?>)
                    Collector.of(supplier, accumulator, combiner, dsCharacteristics);
            return result;
        } else {
            // Finisherを指定する場合
            return Collector.of(supplier, accumulator, combiner,
                                // predicates分のList<A>でFinisherを実行する
                                l -> l.stream().map(downstream.finisher()).collect(Collectors.toList()),
                                dsCharacteristics);
        }
    }

    public static void main(String[] args) {
        List<String> input = asList("abc", "ade", "bcd", "cc", "cdac");

        System.out.println("- 条件１つでtrue/false");
        input.stream()
                .collect(partitioningBy(s -> s.startsWith("a")))
                .entrySet()
                .forEach(System.out::println);


        System.out.println("- 上の例を条件２つで表現");
        input.stream()
                .collect(classifyingBy(
                        Collectors.toSet(),
                        s -> s.startsWith("a"),
                        s -> !s.startsWith("a")
                ))
                .forEach(System.out::println);

        System.out.println("- 条件が２つ（条件を満たせば，両方に分類される）");
        input.stream()
                .collect(classifyingBy(
                        Collectors.toList(),  // <- 戻り値はSetでなくても良い
                        s -> s.length() == 3,
                        s -> s.endsWith("c")
                ))
                .forEach(System.out::println);

        System.out.println("- 条件が３つ（条件を満たさない場合は空になる）");
        input.stream()
                .collect(classifyingBy(
                        Collectors.toSet(),
                        s -> s.length() == 3,
                        s -> s.startsWith("x"),
                        s -> s.endsWith("c")
                ))
                .forEach(System.out::println);

        System.out.println("- 条件を与えなければ，何も分類されない（すべて空）");
        input.stream()
                .collect(classifyingBy(Collectors.toSet()))
                .forEach(System.out::println);

        System.out.println("- collect(Collects.toList())と同じ---");
        input.stream()
                .collect(classifyingBy(Collectors.toSet(),
                                       s -> true))
                .forEach(System.out::println);

    }

}