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

时间:2015-03-26 20:23:06

标签: python arrays numpy sum unique

我有两个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])

数组非常长,因此迭代它(并始终查找相应的伙伴字段所在的位置)效率太低。也许我只是没有看到树木的木材......我希望有人可以帮我解决这个问题!

3 个答案:

答案 0 :(得分:2)

我不知道它是否足够有效(因为我会迭代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])

答案 1 :(得分:1)

使用以下列表理解:

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

您可以使用field获取np.where中相同元素的索引:

>>> [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的相应元素!并使用sum计算np.sum

答案 2 :(得分:0)

使用numpy_indexed包以完全向量化的方式解决此问题(免责声明:我是其作者)

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

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