在 1 个 I/O 通道中提取 numpy 和字典的最快方法

Posted 2023-03-11

【中文标题】在 1 个 I/O 通道中提取 numpy 和字典的最快方法

【英文标题】：Fastest way to extract dictionary of sums in numpy in 1 I/O pass

【发布时间】：2016-11-24 14:02:20

【问题描述】：

假设我有一个像这样的数组：

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 亿行，第一列中有 100k 个唯一元素。

方法 1：天真地，我可以调用 np.unique() 然后在包含列表理解的单行字典中使用 np.where() 和 np.sum() 的组合遍历唯一数组中的每个项目。如果我有少量唯一元素，这将相当快，但是在 100k 唯一元素时，我会浪费大量的页面获取，从而对整个数组进行 100k I/O 传递。

方法 2：我也可以对最后一列进行一次 I/O 传递（因为必须在每一行对第 1 列进行哈希处理可能比过多的页面获取更便宜），但是我失去了 numpy 的 C 内部的优势循环向量化。

有没有一种无需借助纯 Python 循环即可快速实现方法 2 的方法？

【问题讨论】：

何不同时尝试一下，看看效果如何？

@NPE 糟糕，我正在考虑采用方法 2，但我只熟悉在纯 Python 循环中执行此操作，我知道有 @np.vectorize、numba 和 Cython 有助于规避这一点，但也许有一些聪明的方法可以单独使用 numPy 库函数来做到这一点，有人可以指出我...

不确定这会是多快/多慢，但你可以用 pandas groupby 和 sum 来做到这一点

【参考方案1】：

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

【讨论】：

很好的建议，我最终使用了您建议的 pandas 路线，谢谢！

【参考方案2】：

这是一个基于 NumPy 的方法，使用 np.add.reduceat -

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]])

【参考方案3】：

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

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

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

【参考方案4】：

使用 NumPy 的正确方法是使用 np.bincount。如果您唯一的第一列标签已经是小的连续整数，您可以简单地这样做：

cum_sums = np.bincount(arr[:, 0], weights=arr[:, 2])
cum_dict = index: cum_sum for index, cum_sum in enumerate(cum_sums)
            if cum_sum != 0

cum_sum != 0 试图跳过缺少的第一列标签，如果您的第三列包含负数，这可能是严重错误的。

或者你也可以做正确的事，先打电话np.unique然后做：

uniques, indices = np.unique(arr[:, 0], return_inverse=True)
cum_sums = np.bincount(indices, weights=arr[:, 2])
cum_dict = index: cum_sum for index, cum_sum in zip(uniques, cum_sums)