在1 I/O传递中以numpy提取和的字典的最快方法

Question

假设我有一个类似的数组:

arr = np.array([[1,20,5],
                [1,20,8],
                [3,10,4],
                [2,30,6],
                [3,10,5]])

我想形成一个字典,列出第一列的总和,每行匹配第一列中的每个值,即返回{1: 13, 2: 6, 3: 9}.为了使事情更具挑战性,我的数组中有10亿行,第一列中有10万个独特元素.

方法1:天真地,我可以调用np.unique()然后迭代遍历唯一数组中的每个项目,np.where()并np.sum()在一个单行字典中包含列表理解.如果我有少量独特的元素,这将是相当快的,但在100k独特的元素,我将招致大量浪费的页面提取,使整个阵列的100k I/O传递.

方法2:我可以对最后一列进行单个I/O传递(因为必须在每一行中散列第1列可能比过多的页面提取更便宜),但是我在这里失去了numpy的C内循环向量化的优势.

有没有一种快速的方法来实现方法2而不诉诸纯Python循环？

Answer 1

麻木方法：

u = np.unique(arr[:, 0])
s = ((arr[:, [0]] == u) * arr[:, [2]]).sum(0)

dict(np.stack([u, s]).T)

{1: 13, 2: 6, 3: 9}

熊猫方法：

import pandas as pd
import numpy as np

pd.DataFrame(arr, columns=list('ABC')).groupby('A').C.sum().to_dict()

{1: 13, 2: 6, 3: 9}