高效地对NumPy数组中与另一个数组索引匹配的元素进行求和。

Question

我想找出与另一个numpy相匹配的索引对应行的总和。

下面的示例更能说明问题。

A=np.array(['a-1','b-1','b-1','c-2','a-1','b-1','c-2']);
b = np.array([1.21,2.34,1.2,2.8,10.0,0.9,8.4]);;

我希望输出为一个字典，使得：

d['a-1'] = 1.21 + 10.0 = 11.21
d['b-1'] = 2.34 + 1.2 + 0.9 = 4.44
d['c-2'] = 2.8 + 8.4 = 11.2

结果是b数组中与A数组相应的索引处出现相同值的元素之和。有没有一种高效的方法来实现这个功能？我的数组非常大（数百万级别）。

- Shew

Answer 1

方法一

我们可以使用 np.unique 和 np.bincount 的组合 -

In [48]: unq, ids = np.unique(A, return_inverse=True)

In [49]: dict(zip(unq, np.bincount(ids, b)))
Out[49]: 
{'a-1': 11.210000000000001,
 'b-1': 4.4400000000000004,
 'c-2': 11.199999999999999}

因此，np.unique 为 A 中的每个字符串提供了唯一的整数映射，然后将这些整数作为箱子用于基于权重的加权求和，其中权重来源于 b。

方法2（特定情况）

假设 A 中的字符串始终是3个字符，则更快的方法是将这些字符串转换为数字，并将其用作 np.unique 的输入。 这种方法的想法是，np.unique 在数字上的工作比在字符串上的工作快。

因此，实现如下 -

In [141]: n = A.view(np.uint8).reshape(-1,3).dot(256**np.arange(3))

In [142]: unq, st, ids = np.unique(n, return_index=1, return_inverse=1)

In [143]: dict(zip(A[st], np.bincount(ids, b)))
Out[143]: 
{'a-1': 11.210000000000001,
 'b-1': 4.4400000000000004,
 'c-2': 11.199999999999999}

神奇的是，重塑后的 viewing 仍然保持为视图，因此应该非常高效：

In [150]: np.shares_memory(A,A.view(np.uint8).reshape(-1,3))
Out[150]: True

我们可以使用np.unique的axis参数（在1.13.0中添加的功能）

In [160]: A2D = A.view(np.uint8).reshape(-1,3)

In [161]: unq, st, ids = np.unique(A2D, axis=0, return_index=1, return_inverse=1)

In [162]: dict(zip(A[st], np.bincount(ids, b)))
Out[162]: 
{'a-1': 11.210000000000001,
 'b-1': 4.4400000000000004,
 'c-2': 11.199999999999999}

- Divakar

Answer 2

一种替代方案是使用 pandas：

import pandas as pd
df = pd.DataFrame(data=[pd.Series(A),pd.Series(b)]).transpose()
res = df.groupby(0).sum()

提供

res
Out[62]: 
         1
0         
a-1  11.21
b-1   4.44
c-2  11.20

您可以通过以下方式获取所需的字典：

res_dict = res[1].to_dict()

这就是

Out[64]: 
{'a-1': 11.210000000000001,
 'b-1': 4.4400000000000004,
 'c-2': 11.199999999999999}

- GPhilo

Answer 3

numpy_indexed 包（声明：我是它的作者）包含了以高效和优雅的方式执行这些类型操作的功能：

import numpy_indexed as npi
k, v = npi.group_by(A).sum(b)
d = dict(zip(k, v))

我觉得pandas的分组语法有些笨重，执行如此基本的操作不应该需要重新组织数据结构。

- Eelco Hoogendoorn

---

虽然我同意将数据转换为pandas格式的笨拙性，但你可能需要在回答中添加你是该软件包的开发者，否则它可能会被标记为此答案所述。 - GPhilo