com.mogudiandian.util.stream.RandomKCollector Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of joshua-util Show documentation
Joshua's user-friendly utilities
The newest version!
package com.mogudiandian.util.stream;

import javax.annotation.concurrent.NotThreadSafe;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * 实现RandomK的collector
 * 使用方式：
 *  *     int[] arr = new int[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
 *     List list = Arrays.stream(arr).boxed().collect(RandomKCollector.collect(6));
 *     System.out.println(list);
 * 
 *
 * @param  stream中元素的类型
 * @author Joshua Sun
 * @since 1.0.0
 */
@NotThreadSafe
public final class RandomKCollector implements Collector>, List> {

    private int k;

    private RandomKCollector(int k) {
        if (k <= 0) {
            throw new IllegalArgumentException("K must be positive");
        }
        this.k = k;
    }

    /**
     * 创建RandomK收集器实例
     * @param k 需要多少个元素
     * @param  元素类型
     * @return RandomK收集器
     */
    public static  RandomKCollector collect(int k) {
        return new RandomKCollector<>(k);
    }

    @Override
    public Supplier>> supplier() {
        return () -> new PriorityQueue<>(k, Map.Entry.comparingByKey());
    }

    @Override
    public BiConsumer>, T> accumulator() {
        return (queue, t) -> {
            queue.offer(new AbstractMap.SimpleImmutableEntry<>(Math.random(), t));
            if (queue.size() > k) {
                queue.poll();
            }
        };
    }

    @Override
    public BinaryOperator>> combiner() {
        return (q1, q2) -> {
            q1.addAll(q2);
            return q1;
        };
    }

    @Override
    public Function>, List> finisher() {
        return queue -> IntStream.iterate(0, x -> x + 1)
                                 .limit(Math.min(k, queue.size()))
                                 .mapToObj(x -> queue.poll())
                                 .filter(Objects::nonNull)
                                 .map(Map.Entry::getValue)
                                 .collect(Collectors.toList());
    }

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