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