在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内部循环矢量化的优势。你知道吗

有没有一种不依赖于纯Python循环的快速方法来实现方法2？你知道吗

Answer 1

这是一个典型的分组问题，numpy_indexed包高效而优雅地解决了这个问题（如果我自己也这么说的话；我是它的作者）

import numpy_indexed as npi
npi.group_by(arr[:, 0]).sum(arr[:, 2])

它是一个比pandas包更轻量级的解决方案，而且我认为语法更简洁，因为不需要仅仅为了执行这种基本操作而创建特殊的数据结构。性能应该与Divakar提出的解决方案相同，因为它遵循相同的步骤；只是在顶部有一个漂亮且经过测试的接口。你知道吗

Answer 2

numpy方法：

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}

Answer 3

下面是一个基于NumPy的方法，使用^{}-

sidx = arr[:,0].argsort()
idx = np.append(0,np.where(np.diff(arr[sidx,0])!=0)[0]+1)
keys = arr[sidx[idx],0]
vals = np.add.reduceat(arr[sidx,2],idx,axis=0)

如果要获取2列数组中的键和值-

out = np.column_stack((keys,vals)) # If you 

样本运行-

In [351]: arr
Out[351]: 
array([[ 1, 20,  5],
       [ 1, 20,  8],
       [ 3, 10,  4],
       [ 2, 30,  6],
       [ 3, 10,  5]])

In [352]: out
Out[352]: 
array([[ 1, 13],
       [ 2,  6],
       [ 3,  9]])