Python按数组a分组并汇总数组b - 性能

Question

给定两个长度相同的无序数组 a 和 b：

a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]

我想根据数组 a 中的元素进行分组：

aResult = [7,3,5]

并对数组 b 中的元素进行求和（这个例子是用来总结概率密度函数的）：

bResult = [0.2 + 0.1 + 0.2, 0.1, 0.3 + 0.1] = [0.5, 0.1, 0.4]

另外，可以在 Python 中随机生成 a 和 b：

import numpy as np
a = np.random.randint(1,10,10000)
b = np.array([1./len(a)]*len(a))

我有两种方法，肯定还没有达到性能的最低边界。方法 1（至少简单又短）：耗时：0.769315958023

def approach_2(a,b):
    bResult = [sum(b[i == a]) for i in np.unique(a)]
    aResult = np.unique(a)

方法 2（使用 numpy.groupby，速度非常慢）耗时：4.65299129486

def approach_2(a,b): 
    tmp = [(a[i],b[i]) for i in range(len(a))]
    tmp2 = np.array(tmp, dtype = [('a', float),('b', float)])
    tmp2 = np.sort(tmp2, order='a') 

    bResult = []
    aResult = []
    for key, group in groupby(tmp2, lambda x: x[0]):
        aResult.append(key)
        bResult.append(sum([i[1] for i in group]))

更新：方法 3，由 Pablo 提供。耗时：1.0265750885

def approach_Pablo(a,b):    

    pdf = defaultdict(int); 
    for x,y in zip(a,b):
        pdf[x] += y  

更新：方法 4，由 Unutbu 提供。耗时：0.184849023819 [到目前为止的赢家，但 a 只能是整数]

def unique_Unutbu(a,b):

    x=np.bincount(a,weights=b)
    aResult = np.unique(a)
    bResult = x[aResult]

也许有人能找到比我更聪明的解决方案 :)

Answer 1

这个方法怎么样：

from collections import defaultdict
pdf = defaultdict(int)
a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]
for x,y in zip(a,b):
  pdf[x] += y

你只需要遍历每个元素一次，并使用字典来快速查找。如果你真的想要最后得到两个单独的数组，你可以请求它们：

aResult = pdf.keys()
bResult = pdf.values()

Answer 2

这里有一种方法，和@unutbu的类似：

import numpy as np

def f(a, b):
    result_a, inv_ndx = np.unique(a, return_inverse=True)
    result_b = np.bincount(inv_ndx, weights=b)
    return result_a, result_b

这个方法允许数组使用非整数类型，也就是说，你可以用其他类型的数据。它还支持数组中有很大的数值。返回的结果是数组的元素按顺序排列的。如果你想保留原来的顺序，可以使用np.unique()函数中的return_index参数来做到。

不过，随着数组中独特元素的数量增加，这个方法的性能会变得更差。在你的问题中，这个方法的速度比@unutbu的方法慢了四倍。

我做了一个性能对比，还加入了另外三种方法。表现最好的方法是：对于整数数组，使用基于哈希的实现，这是用Cython写的；对于double数组（输入大小为10000），表现最好的方法是基于排序的实现，同样是用Cython写的。

Answer 3

如果 a 里面的数字都是小于 2³¹-1 的整数（也就是说，这些数字可以放进 int32 这种数据类型里），那么你可以用 np.bincount 来计算加权值：

import numpy as np
a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]

x=np.bincount(a,weights=b)
print(x)
# [ 0.   0.   0.   0.1  0.   0.4  0.   0.5]

print(x[[7,3,5]])
# [ 0.5  0.1  0.4]

np.unique(a) 会返回 [3 5 7]，所以结果的顺序可能会不一样：

print(x[np.unique(a)])
# [ 0.1  0.4  0.5]

使用 np.bincount 可能会遇到一个问题，就是它返回的数组长度会等于 a 中的最大值。如果 a 里面有一个数字接近 2³¹-1，那么 bincount 就得分配一个大小为 8*(2**31-1) 字节（也就是 16GiB）的数组。

所以，对于长度很大的数组 a，如果里面的值不大，使用 np.bincount 可能是最快的解决方案。但如果数组 a 的长度比较小（不管值是大还是小），那么使用 collections.defaultdict 可能会更快。

补充：可以查看 J.F. Sebastian 的解决方案，它提供了一种绕过只使用整数值和大值问题的方法。