大量数据topk-分桶+堆+多路并归解决方案

利用分桶、堆与多路归并解决 TopK 问题:结果处理阶段解析

在处理大规模数据时,TopK 问题是一个常见且具有挑战性的任务,即从海量数据中找出最大(或最小)的 K 个元素。为了高效地解决这个问题,我们可以采用分桶、堆和多路归并相结合的方法。本文将详细剖析该方法中结果处理阶段的代码逻辑。

问题背景

TopK 问题在数据处理、搜索引擎、推荐系统等领域都有广泛的应用。为了高效解决该问题,我们采用了分桶、堆和多路归并的策略。具体步骤包括:首先将数据分桶,降低数据规模;然后在每个桶中使用最小堆找出局部的 TopK 元素;最后将每个桶的 TopK 元素合并到全局最小堆中

具体代码

复制代码
`import java.util.*;

public class TopKSolution {

    public static List<Integer> topK(int[] nums, int k) {
        // 步骤 1: 分桶
        int min = Integer.MAX_VALUE;
        int max = Integer.MIN_VALUE;
        for (int num : nums) {
            min = Math.min(min, num);
            max = Math.max(max, num);
        }
        // 桶的数量
        int bucketSize = 10; 
        int bucketCount = (max - min) / bucketSize + 1;
        List<List<Integer>> buckets = new ArrayList<>();
        for (int i = 0; i < bucketCount; i++) {
            buckets.add(new ArrayList<>());
        }
        // 将元素放入对应的桶中
        for (int num : nums) {
            int bucketIndex = (num - min) / bucketSize;
            buckets.get(bucketIndex).add(num);
        }

        // 步骤 2: 每个桶中使用最小堆找出 TopK
        PriorityQueue<Integer> globalHeap = new PriorityQueue<>(k);
        for (List<Integer> bucket : buckets) {
            if (bucket.isEmpty()) continue;
            // 为当前桶创建一个容量为 k 的最小堆,用于找出该桶内的 TopK 元素
            PriorityQueue<Integer> localHeap = new PriorityQueue<>(k, Comparator.naturalOrder());
            for (int num : bucket) {
                if (localHeap.size() < k) {
                    localHeap.offer(num);
                } else if (num > localHeap.peek()) {
                    localHeap.poll();
                    localHeap.offer(num);
                }
            }
            // 将每个桶的 TopK 元素合并到全局堆中
            for (int num : localHeap) {
                if (globalHeap.size() < k) {
                    globalHeap.offer(num);
                } else if (num > globalHeap.peek()) {
                    globalHeap.poll();
                    globalHeap.offer(num);
                }
            }
        }

        // 步骤 3: 结果处理
        List<Integer> result = new ArrayList<>(globalHeap);
        result.sort(Collections.reverseOrder());
        return result;
    }

    public static void main(String[] args) {
        int[] nums = {3, 2, 1, 5, 6, 4};
        int k = 2;
        List<Integer> topK = topK(nums, k);
        System.out.println("Top " + k + " elements: " + topK);
    }
}    `

代码解释

分桶:

  1. 先找出数组里的最小值 min 和最大值 max。
  2. 确定桶的数量 bucketCount,这里每个桶的大小为 bucketSize。
  3. 把数组中的每个元素依据其值放入对应的桶中。

堆:

  1. 针对每个桶,使用最小堆 localHeap 找出该桶内的 TopK 元素。
  2. 要是堆的大小小于 K,就直接将元素加入堆;若堆的大小已达到 K 且当前元素比堆顶元素大,就移除堆顶元素并将当前元素加入堆。

多路归并:

  1. 把每个桶的 TopK 元素合并到全局最小堆 globalHeap 中。
  2. 最终从全局堆中获取最大的 K 个元素。

结果处理:

  1. 把全局堆中的元素存到列表里,然后按降序排序。