我可以限制Collectors.toMap()条目吗?

时间:2016-09-02 22:09:01

标签: java java-stream collectors

我正在寻找一种方法来限制Collectors.toMap()使用合并功能生成的条目数。请考虑以下示例:

Map<String, Integer> m = Stream.of("a", "a", "b", "c", "d")
    .limit(3)
    .collect(toMap(Function.identity(), s -> 1, Integer::sum));

上面的问题是我在结果地图中只有2个元素(a=2, b=1)。有没有方便的方法在流处理3 不同的键时使流短路?

1 个答案:

答案 0 :(得分:2)

一个可能的解决方案是编写自己的Spliterator,它将包装给定Stream的分裂器。这个Spliterator会将前进的调用委托给包装的spliterator并包含计数许多不同元素的逻辑。

为此,我们可以继承AbstractSpliterator并提供我们自己的tryAdvance逻辑。在下文中,遇到的所有元素都添加到集合中。当该集合的大小变得大于我们的最大值或者当包装的分裂器没有剩余元素时,我们返回false以指示没有剩余的元素需要考虑。当达到不同元素的数量时,这将停止。

private static <T> Stream<T> distinctLimit(Stream<T> stream, int max) {
    Spliterator<T> spltr = stream.spliterator();
    Spliterator<T> res = new AbstractSpliterator<T>(spltr.estimateSize(), spltr.characteristics()) {

        private Set<T> distincts = new HashSet<>();
        private boolean stillGoing = true;

        @Override
        public boolean tryAdvance(Consumer<? super T> action) {
            boolean hasRemaining = spltr.tryAdvance(elem -> {
                distincts.add(elem);
                if (distincts.size() > max) {
                    stillGoing = false;
                } else {
                    action.accept(elem);
                }
            });
            return hasRemaining && stillGoing;
        }
    };
    return StreamSupport.stream(res, stream.isParallel()).onClose(stream::close);
}

使用您的示例代码,您将拥有:

Map<String, Long> m =
    distinctLimit(Stream.of("a", "a", "b", "c", "d"), 3)
        .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));

并且输出将是预期的{a=2, b=1, c=1},即具有3个不同键的地图。