NumbaPro - 对二维数组进行排序然后对相同键的条目求和的最聪明方法

Question

在我的程序中，我有一个大小为数百万个条目的数组，如下所示：

arr=[(1,0.5), (4,0.2), (321, 0.01), (2, 0.042), (1, 0.01), ...]

如果有帮助，我可以改为使用相同的顺序制作两个数组（而不是带有元组的数组）。

为了对这个数组进行排序，我知道我可以使用基数排序，因此它具有以下结构：

arr_sorted = [(1,0.5), (1,0.01), (2,0.42), ...]

现在我想对数组中具有键 1 的所有值求和。然后所有具有键 2 等的值应该写入一个新数组，如下所示：

arr_summed = [(1, 0.51), (2,0.42), ...]

显然这个数组会小得多，尽管仍然在 100000 个条目的数量级。 现在我的问题是：在 CUDA 中解决我的问题的最佳并行方法是什么？ 我正在使用 NumbaPro。

为清楚起见进行编辑

我会有两个数组而不是像这样的元组列表：

keys = [1, 2, 5, 2, 6, 4, 4, 65, 3215, 1, .....]
values = [0.1, 0.4, 0.123, 0.01, 0.23, 0.1, 0.1, 0.4 ...]

它们最初是被复制到设备的 numpy 数组。

我想要的是按键减少它们，如果可能的话，将缺失的键值（例如，如果三个未出现在数组中）设置为零。

所以我希望它变成：

keys = [1, 2, 3, 4, 5, 6, 7, 8, ...]
values = [0.11, 0.41, 0, 0.2, ...] # <- Summed by key

我事先知道最终数组有多大。

Answer 1

我不知道 Numba，但在简单的 Python 中：

arr=[(1,0.5), (4,0.2), (321, 0.01), (2, 0.042), (1, 0.01), ...]
res = [0.0] * (indexmax + 1)
for k, v in arr:
   res[k] += v