Python group by array a,并汇总数组b - Performance

Question

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

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

我想按以下要素分组:

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,由巴勃罗.时间:1.0265750885

def approach_Pablo(a,b):    

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

更新:Unutbu的方法4.时间:0.184849023819 [WINNER SO FAR,但仅限整数]

def unique_Unutbu(a,b):

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

也许有人找到比我更聪明的解决方案来解决这个问题:)

Answer 1

这里的方法类似于@ 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

它允许非整数类型的a数组.它允许a数组中的大值.它a以排序顺序返回元素.如果需要,可以使用return_index函数参数轻松恢复原始顺序np.unique().

随着独特元素数量的a增加,性能会恶化.它比你问题的数据慢了4倍于@ unutbu的版本.

我用其他三种方法进行了性能比较.领导者是:对于整数数组 - 在Cython中基于哈希的实现 ; 对于double数组(对于输入大小10000) - 基于排序的impl.也在Cython中.

Answer 2

如果a由int <2**31-1组成(即,如果a具有可以适合dtype的值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**31-1的元素,则bincount必须分配一个大小8*(2**31-1)字节数组(或16GiB).

因此,np.bincount对于a具有较大长度但不是较大值的阵列,可能是最快的解决方案.对于a具有较小长度(以及大或小值)的数组,使用a collections.defaultdict可能会更快.

编辑:请参阅JF Sebastian的解决方案,以解决仅限整数值限制和大值问题的方法.