Numpy如何对数字和数组求和?

内容来源于 Stack Overflow,并遵循CC BY-SA 3.0许可协议进行翻译与使用

  • 回答 (2)
  • 关注 (0)
  • 查看 (76)

假设我有一个Numpy数组,如下所示:1,2,3,4,5,6另一个数组:0,0,1,2,2,1我希望按组(第二个数组)对第一个数组中的项进行求和,并以组号顺序获得n组结果(在本例中,结果将是3、9、9)。我该怎么做?

提问于
用户回答回答于

有不止一种方法可以做到这一点,但这里有一种方法:

import numpy as np
data = np.arange(1, 7)
groups = np.array([0,0,1,2,2,1])

unique_groups = np.unique(groups)
sums = []
for group in unique_groups:
    sums.append(data[groups == group].sum())

向量化,这样就根本没有for循环了,但我建议您不要这样做。它变得不可读,并且需要两个2D临时数组,如果您有大量的数据,这可能需要大量的内存。

import numpy as np
data = np.arange(1, 7)
groups = np.array([0,0,1,2,2,1])

unique_groups = np.unique(groups)

# Forgive the bad naming here...
# I can't think of more descriptive variable names at the moment...
x, y = np.meshgrid(groups, unique_groups)
data_stack = np.tile(data, (unique_groups.size, 1))

data_in_group = np.zeros_like(data_stack)
data_in_group[x==y] = data_stack[x==y]

sums = data_in_group.sum(axis=1)
用户回答回答于

这是一种矢量化的方法,它基于numpy.unique的实现来完成这个和。根据我的时间表,它比循环法快500倍,比直方图法快100倍。

def sum_by_group(values, groups):
    order = np.argsort(groups)
    groups = groups[order]
    values = values[order]
    values.cumsum(out=values)
    index = np.ones(len(groups), 'bool')
    index[:-1] = groups[1:] != groups[:-1]
    values = values[index]
    groups = groups[index]
    values[1:] = values[1:] - values[:-1]
    return values, groups

扫码关注云+社区