Top-N 泛型工具类

呦，又写BUG呢2024-02-25 10:28

一、代码实现

通过封装 PriorityQueue 实现，PriorityQueue 本质上是完全二叉树实现的小根堆（相对来说，如果比较器反向比较则是大根堆）。

java 复制代码

public class TopNUtil<E extends Comparable<E>> {
    private final PriorityQueue<E> priorityQueue;
    private final int n;

    /**
     * 构造 Top-N
     */
    public TopNUtil(int size) {
        if (size <= 0) {
            throw new IllegalArgumentException("Top-N size must be a positive number");
        }
        this.priorityQueue = new PriorityQueue<>(size);
        this.n = size;
    }

    /**
     * 向 Top-N 中插入元素
     */
    public void add(E e) {
        if (priorityQueue.size() < n) {
            priorityQueue.add(e);
            return;
        }
        
        E head = priorityQueue.peek();
        if (head != null && e.compareTo(head) <= 0) {
            return;
        }
        
        priorityQueue.poll();
        priorityQueue.add(e);
    }

    /**
     * 将 Top-N 转为从大到小排序的 List
     */
    public List<E> toSortedArrayList() {
        List<E> tempList = new ArrayList<>(priorityQueue);
        tempList.sort(Collections.reverseOrder());
        return tempList;
    }
}

二、使用示例

java 复制代码

class TopNUtilTest {

    @Test
    void test() {
        List<CountDTO> list = new ArrayList<>();
        TopNUtil<CountDTO> top = new TopNUtil<>(3);

        // 生成 10 个随机的 CountDTO
        for (int i = 0; i < 10; i++) {
            CountDTO dto = new CountDTO();
            dto.setOrderPriceSum(BigDecimal.valueOf(Math.random() * 100));
            list.add(dto);
        }

        System.out.println("所有的 CountDTO 值:");
        for (CountDTO dto : list) {
            System.out.print(dto.getOrderPriceSum());
            System.out.print(" ");
        }
        System.out.println();

        // 将列表中的元素添加到 TopNUtil
        for (CountDTO dto : list) {
            top.add(dto);
        }

        // 获取 TopNUtil 中的元素列表
        List<CountDTO> topList = top.toSortedArrayList();

        // 确保列表的大小不超过 3
        assertEquals(3, topList.size());

        // 打印 Top 3 元素的 CountDTO 值
        System.out.println("Top 3 CountDTO 值:");
        for (CountDTO dto : topList) {
            System.out.println(dto.getOrderPriceSum());
        }
    }
}