使用NumPy按数字对数组进行求和

30

假设我有一个像这样的numpy数组: [1,2,3,4,5,6] 还有另一个数组: [0,0,1,2,2,1] 我想按组(第二个数组)对第一个数组中的项目进行求和,并按组号顺序获得n组结果(在这种情况下,结果将是[3, 9, 9])。我该如何在numpy中实现此操作?


你为什么需要numpy呢?难道你不只是使用原生的Python列表吗?如果不是,那你使用的是哪种numpy类型? - Matt Ball
2
我需要使用numpy,因为我不想为n个组循环遍历数组n次,因为我的数组大小可能是任意大的。我没有使用Python列表,只是在括号中展示了一个示例数据集。数据类型是int。 - Scribble Master
与https://dev59.com/U1rUa4cB1Zd3GeqPi1iB相关的编程内容。 - TooTone
11个回答

57

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 元素的总和。

希望这可以帮到您。


1
可以确认这非常快。对于小输入,大约比Bi Rico提供的sum_by_group方法快10倍。 - Jonathan Richards
如果“data”是向量怎么办? - zzh1996
看起来权重参数必须是一维的。一个解决方案是为向量的每个维度运行一次bincount(如果数据是一组二维向量,则运行两次)。对Peter答案的轻微修改也可以起作用。 - Alex
2
很棒的方法。请注意,bincount需要“int”类型的id。 - Mark_Anderson
请注意,即使您期望出现的所有ID都不存在,这也是有意义的。 - Simon P
有人找到了float类型id的解决方法吗? - David Cian

30

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

14

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

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

您可以将事物向量化,以便根本没有循环,但我建议不要这样做。它变得难以阅读,并且需要几个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)

@Scribble Master - 看看修改...循环唯一组没有问题。第二个版本可能会很慢,而且很难读懂。使用循环时,你只需要在(至少在Python中)循环唯一组的数量上进行循环。内部比较data[groups == group]将非常快速。 - Joe Kington
这个 data[groups == group] 的语法是什么黑魔法?将一个数组与标量进行比较会产生某种切片或视图?o_O - Karl Knechtel
@Karl - groups == group 会产生一个布尔数组。在numpy中,你可以通过数组进行索引。这是numpy(和Matlab)中非常常见的用法。我认为它非常易读(把它看作“where”),而且非常有用。 - Joe Kington
@Joe:很整洁,但对我来说可能有点太神奇了。我没有做过太多关于Numpy的事情(没有像我想象中那么需要它)-这需要一些时间来适应。 - Karl Knechtel
@JoeKington 很棒的回答!快速而且高效。 - Canopus
显示剩余2条评论

7
如果分组是按连续整数索引的,您可以滥用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循环。

7

我尝试了其他人的脚本,我的考虑如下:

用户 评论
乔 (Joe) 只适用于少数组。
凯文 (kevpie) 由于循环太慢(这不是Python编程的方式)。
Bi_Rico和斯文 (Sven) 性能不错,但仅适用于Int32 (如果总和超过2^32/2,则会失败)。
亚历克斯 (Alex) 最快的方法,对于求和来说是最佳解决方案。

但如果您想要更多的灵活性和按其他统计数据分组的可能性,请使用SciPy

import numpy as np
from scipy import ndimage

data = np.arange(10000000)
unique_groups = np.arange(1000)
groups = unique_groups.repeat(10000)

ndimage.sum(data, groups, unique_groups)

这是很好的,因为你可以有许多统计数据进行分组(求和、平均数、方差等)。


1
这个解决方案非常简洁。 - notilas

5

你们都错了!做这件事最好的方法是:

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.])

实际上,你应该只是使用Alex的np.bincount答案。 - Peter

2

我注意到了 numpy 标签,但如果你不介意使用 pandas,这个任务可以变成一行代码:

import pandas as pd
import numpy as np

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

df = pd.DataFrame({'data': data, 'groups': groups})

所以 df 的样子是这样的:
   data  groups
0     1       0
1     2       0
2     3       1
3     4       2
4     5       2
5     6       1

现在你可以使用函数groupby()sum()

print(df.groupby(['groups'], sort=False).sum())

这将为您提供所需的输出结果。

        data
groups      
0          3
1          9
2          9

默认情况下,数据框会被排序,因此我使用了标志sort=False,这可能会提高大型数据框的速度。


1

我尝试了不同的方法来做这件事,发现使用np.bincount确实是最快的。请参见Alex的答案。

    import numpy as np
    import random
    import time
    
    size = 10000
    ngroups = 10
    
    groups = np.random.randint(low=0,high=ngroups,size=size)
    values = np.random.rand(size)
    
    
    # Test 1                                                                                                                                                                                                           
    beg = time.time()
    result = np.zeros(ngroups)
    for i in range(size):
        result[groups[i]] += values[i]
    print('Test 1 took:',time.time()-beg)
    
    # Test 2                                                                                                                                                                                                           
    beg = time.time()
    result = np.zeros(ngroups)
    for g,v in zip(groups,values):
        result[g] += v
    print('Test 2 took:',time.time()-beg)
    
    # Test 3                                                                                                                                                                                                           
    beg = time.time()
    result = np.zeros(ngroups)
    for g in np.unique(groups):
        wh = np.where(groups == g)
        result[g] = np.sum(values[wh[0]])
    print('Test 3 took:',time.time()-beg)
    
    
    # Test 4                                                                                                                                                                                                           
    beg = time.time()
    result = np.zeros(ngroups)
    for g in np.unique(groups):
        wh = groups == g
        result[g] = np.sum(values, where = wh)
    print('Test 4 took:',time.time()-beg)
    
    # Test 5                                                                                                                                                                                                           
    beg = time.time()
    result = np.array([np.sum(values[np.where(groups == g)[0]]) for g in np.unique(groups) ])
    print('Test 5 took:',time.time()-beg)
    
    # Test 6                                                                                                                                                                                                           
    beg = time.time()
    result = np.array([np.sum(values, where = groups == g) for g in np.unique(groups) ])
    print('Test 6 took:',time.time()-beg)
    
    # Test 7                                                                                                                                                                                                           
    beg = time.time()
    result = np.bincount(groups, weights = values)
    print('Test 7 took:',time.time()-beg)

结果:

    Test 1 took: 0.005615234375
    Test 2 took: 0.004812002182006836
    Test 3 took: 0.0006084442138671875
    Test 4 took: 0.0005099773406982422
    Test 5 took: 0.000499725341796875
    Test 6 took: 0.0004980564117431641
    Test 7 took: 1.9073486328125e-05

1

另外,注意Alex的回答:

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

如果你的索引不是连续的,你可能会陷入困境,想知道为什么一直得到许多零值。
例如:
data = [1,2,3,4,5,6]
ids  = [1,1,3,5,5,3]
np.bincount(ids, weights=data)

会给你:
array([0, 3, 0, 9, 0, 9])

这显然意味着它会从列表中的0到max id构建所有唯一的箱子,然后返回每个箱子的总和。

0
这里有一个方法,可以对任何维度的对象进行分组求和,而且可以根据任何类型的值(不仅限于int)进行分组:
grouping = np.array([1.1, 10, 1.1, 15])
to_sum = np.array([
    [1, 0],
    [0, 1],
    [0.5, 0.3],
    [2, 5],
])

groups, element_group_ixs = np.unique(grouping, return_inverse=True)
accum = np.zeros((groups.shape[0], *to_sum.shape[1:]))
np.add.at(accum, element_group_ixs, to_sum)

结果是:

groups = array([ 1.1, 10. , 15. ])
accum = array([
    [1.5, 0.3],
    [0. , 1. ],
    [2. , 5. ]
])

(np.add.at的想法来自Peter的回答)


网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接