Top K Selection

原始题目：LeetGPU - Top K Selection

题目描述

编写一个 GPU 程序，给定一个长度为 $N$ 的 32 位浮点数一维数组 input，选出最大的 $k$ 个元素，按降序写入长度为 $k$ 的输出数组。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 数组中。

示例

示例 1

Input:  input = [1.0, 5.0, 3.0, 2.0, 4.0], N = 5, k = 3
Output: [5.0, 4.0, 3.0]

示例 2

Input:  input = [7.2, -1.0, 3.3, 8.8, 2.2], N = 5, k = 2
Output: [8.8, 7.2]

约束条件

• $1 \le N \le 100{,}000{,}000$，$1 \le k \le N$。
• 所有值为 32 位浮点数。
• 性能测试在 $N = 50{,}000{,}000,\ k = 100$ 的规模下进行。

解题思路

Top-K 是 GPU 上的经典问题。当 $k$ 很小时（如 100），不需要完整排序——使用基于堆的选择或基数选择（Radix Select）更高效。典型做法是：每个 block 用共享内存维护一个大小为 $k$ 的最小堆，扫描自己负责的数据段，最终合并各 block 的 top-k。也可以先用桶排序思想按最高几位分桶，只在候选桶中继续筛选。注意题目要求结果降序，最后一步需要对选出的 $k$ 个元素排序。欢迎在 GitHub Discussions 分享你的解法。