Python/Numpy - 计算相等数组元素的总和

Question

我有两个 numpy 数组，看起来像：

field = np.array([5,1,3,3,2,1,6])    
counts = np.array([100,210,300,150,20,90,170])

它们没有排序（也不应该改变）。我现在想计算第三个数组（具有相同的长度和顺序），只要它们位于同一字段中，它就包含计数的总和。这里的结果应该是：

field_counts = np.array([100,300,450,450,20,300,170])

数组非常长，因此遍历它（并且总是查看相应的合作伙伴字段在哪里）效率太低了。也许我只是没有看到树木的树木......我希望有人可以帮助我！

Answer 1

我不知道它是否足够高效（因为我确实迭代了field）但这里有一个建议。我首先创建一个字段/计数值目录。然后，我根据它创建一个数组。

from collections import defaultdict
dic = defaultdict(int)
for j, f in enumerate(field):
    dic[f] += counts[j]

field_counts = np.array([dic[f] for f in field])

Answer 2

使用以下列表理解：

>>> [np.sum(counts[np.where(field==i)]) for i in field]
[100, 300, 450, 450, 20, 300, 170]

您可以通过np.where 获取field 中相同元素的索引：

>>> [np.where(field==i) for i in field]
[(array([0]),), (array([1, 5]),), (array([2, 3]),), (array([2, 3]),), (array([4]),), (array([1, 5]),), (array([6]),)]

然后通过索引获取counts的对应元素！并用np.sum 计算sum。

Answer 3

这个问题可以使用numpy_indexed 包以完全矢量化的方式解决（免责声明：我是它的作者）

import numpy_indexed as npi
g = npi.group_by(field)
field_counts = g.sum(counts)[1][g.inverse]

g.sum 计算每组唯一字段的总和，然后 g.inverse 将这些值映射回原始字段。