在numpy中按数字求和数组

Question

假设我有一个numpy数组,如:[1,2,3,4,5,6]和另一个数组:[0,0,1,2,2,1]我想按组对第一个数组中的项求和(第二个数组)并获得n组结果的组号顺序(在这种情况下,结果将是[3,9,9]).我怎么做numpy？

Answer 1

numpy函数bincount完全是为了这个目的而制作的,我相信它会比其他方法对所有大小的输入快得多:

data = [1,2,3,4,5,6]
ids  = [0,0,1,2,2,1]

np.bincount(ids, weights=data) #returns [3,9,9] as a float64 array

输出的第i个元素是data对应于"id" 的所有元素的总和i.

希望有所帮助.

Answer 2

这是一种基于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

Answer 3

有多种方法可以做到这一点,但这是一种方式:

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)

Answer 4

如果组由连续的整数索引,则可以滥用该numpy.histogram()函数来获得结果:

data = numpy.arange(1, 7)
groups = numpy.array([0,0,1,2,2,1])
sums = numpy.histogram(groups, 
                       bins=numpy.arange(groups.min(), groups.max()+2), 
                       weights=data)[0]
# array([3, 9, 9])

这将避免任何Python循环.

Answer 5

我尝试过每个人的脚本,我的考虑是:

乔:只有少数几个小组才能工作.

kevpie:由于循环太慢(这不是pythonic方式)

Bi_Rico和Sven:表现不错,但只适用于Int32(如果总和超过2 ^ 32/2则会失败)

亚历克斯:是最快的一个,总和很好.

但是,如果您想要更灵活,并且可以通过其他统计数据进行分组,请使用SciPy:

from scipy import ndimage

data = np.arange(10000000)
groups = np.arange(1000).repeat(10000)
ndimage.sum(data, groups, range(1000))

这很好,因为你有很多分组统计(总和,均值,方差,......).

Answer 6

你们都错了！最好的方法是：

a = [1,2,3,4,5,6]
ix = [0,0,1,2,2,1]
accum = np.zeros(np.max(ix)+1)
np.add.at(accum, ix, a)
print accum
> array([ 3.,  9.,  9.])