package com.koala.common.utils;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * 求topK问题的工具类
 *
 * @author yanghang
 */
public class TopKHeap<T> {

    /**
     * 优先级队列,底层实现类
     */
    PriorityQueue<T> priorityQueue;
    /**
     * 求top几
     */
    int k;
    /**
     * 比较器
     */
    Comparator<T> comparator;

    /**
     * 构造函数
     *
     * @param comparator 比较器
     * @param k          top几
     */
    public TopKHeap(Comparator<T> comparator, int k) {
        this.k = k;
        this.comparator = comparator;
        priorityQueue = new PriorityQueue<>(k, comparator);
    }

    /**
     * 新增元素,底层基于 {@link PriorityQueue#offer(Object)}
     *
     * @param t 需要新增的元素
     * @return true {@link PriorityQueue#offer(Object)}
     */
    public boolean offer(T t) {
        if (priorityQueue.size() < k) {
            priorityQueue.offer(t);
        } else if (comparator.compare(priorityQueue.peek(), t) < 0) {
            priorityQueue.poll();
            priorityQueue.offer(t);
        }
        return true;
    }

    /**
     * 获取topK的集合
     *
     * @return {@link LinkedList}
     */
    public List<T> getTopK() {
        int size = priorityQueue.size();
        List<T> res = new LinkedList<>();
        while (size-- > 0) {
            res.add(priorityQueue.poll());
        }
        return res;
    }

    public static void main(String[] args) {
        TopKHeap<Integer> heap = new TopKHeap<>(Comparator.comparingInt(o -> o), 5);
        Random random = new Random();
        for (int i = 100; i > 0; i--) {
            heap.offer(random.nextInt(100));
        }
        System.out.println(heap.getTopK());
    }

}