用numpy按数字求数组和

Question

假设我有一个numpy数组，内容是：
[1,2,3,4,5,6]
还有另一个数组：
[0,0,1,2,2,1]
我想根据第二个数组中的组别，把第一个数组里的数字加起来，最后得到按组别顺序排列的结果（在这个例子中，结果会是 [3, 9, 9]）。我该如何在numpy中做到这一点呢？

Answer 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

Answer 2

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个元素是所有与“id”i对应的data元素的总和。

希望这能帮到你。

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())

你可以把事情向量化，这样就完全不需要循环了，但我不推荐这样做。这样会让代码变得难以阅读，而且可能需要几个二维的临时数组，如果你的数据量很大，可能会占用大量内存。

补充一下：这里有一种完全向量化的方法。请记住，这种方法可能会比上面提到的版本慢（而且很可能会慢）。(也许还有更好的向量化方法，但现在很晚了，我有点累，所以这只是我脑海中冒出来的第一个想法...)

不过，要记住这是个糟糕的例子... 实际上，使用上面的循环在速度和可读性上都更好...

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)