Java如何基于条件从分组PriorityQueue中按规则取出topN数据

原有代码问题 #

1. 语法错误：static final修饰的比较器不能定义在方法内部，编译无法通过
2. 空间浪费：为每个分组构建全量PriorityQueue存储所有元素，但你只需要每个分组的最高分元素，不需要存储分组内所有数据
3. 时间冗余：全量元素入PriorityQueue的时间复杂度为O(总元素数 * log 单分组元素数)，完全无必要；重复写了两次遍历Map取队首的逻辑，代码冗余
4. 筛选效率待提升：分组数远大于n时，全量排序所有分组最高分的效率低于用固定大小小顶堆筛选的效率
5. 输出顺序不符合示例：原有实现分组数<=n时返回的列表是无序的，和你给出的示例输出顺序不一致

优化实现方案 #

优化后时间复杂度降到O(总元素数 + 分组数 * log n)，空间复杂度降到O(分组数)，远优于原有实现：

import java.util.*;
import java.util.stream.Collectors;

public class DataService {
    // 静态比较器移到方法外部，按score倒序
    private static final Comparator<Data> SCORE_DESC_COMPARATOR = Comparator.comparingDouble(Data::getScore).reversed();

    public List<Data> getTopN(final List<Data> dataList, final int n) {
        // 第一步：只保留每个分组的最高分元素，不需要构建全量PriorityQueue
        Map<String, Data> groupMaxMap = new HashMap<>();
        for (Data data : dataList) {
            String group = data.getGroup();
            if (group == null) {
                continue;
            }
            // 不存在当前分组，或者当前data比已有分组的最高分更高，就替换
            groupMaxMap.compute(group, (k, existingVal) -> 
                existingVal == null || data.getScore() > existingVal.getScore() ? data : existingVal
            );
        }

        // 第二步：统一收集所有分组的最高分
        Collection<Data> allGroupMax = groupMaxMap.values();
        if (allGroupMax.size() <= n) {
            // 分组数<=n直接按score倒序返回，和示例输出格式统一
            return allGroupMax.stream()
                    .sorted(SCORE_DESC_COMPARATOR)
                    .collect(Collectors.toList());
        }

        // 第三步：分组数>n时，用大小为n的小顶堆筛选前n个最高分，比全量排序效率更高
        PriorityQueue<Data> minHeap = new PriorityQueue<>(Comparator.comparingDouble(Data::getScore));
        for (Data data : allGroupMax) {
            if (minHeap.size() < n) {
                minHeap.offer(data);
            } else if (data.getScore() > minHeap.peek().getScore()) {
                minHeap.poll();
                minHeap.offer(data);
            }
        }

        // 小顶堆输出是升序，调整为倒序返回
        List<Data> result = new ArrayList<>(minHeap);
        result.sort(SCORE_DESC_COMPARATOR);
        return result;
    }
}

// 原Data类保持不变即可

优化说明 #

• 分组阶段直接用HashMap存每个分组的最大值，避免了所有元素入队的额外开销，元素量越大性能提升越明显
• 分组数大于n时，用大小为n的小顶堆筛选，当分组数远大于n时（比如有10000个分组），全量排序需要O(10000 log 10000)操作，用小顶堆只需要O(10000 log 3)，性能差距非常大
• 代码逻辑更简洁，没有重复的遍历逻辑，也修正了原有语法错误
• 统一对输出结果按score倒序排序，完全符合你给出的示例输出规则

内容的提问来源于stack exchange，提问作者user3407267