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

I有两种方法，这肯定远离较低的性能边界。 方法 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 [到目前为止的获胜者，但仅作为整数]

def unique_Unutbu(a,b):

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

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

Answer 1

这是类似于 @unutbu 的方法< /a>：

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 元素。如果需要，可以使用 np.unique() 函数的 return_index 参数轻松恢复原始顺序。

随着 a 中唯一元素数量的增加，其性能会恶化。对于您问题中的数据，它比 @unutbu 的版本慢 4 倍。

我与另外三种方法进行了性能比较。领先者是： 对于整数数组 - Cython 中基于哈希的实现；对于 double 数组（对于输入大小 10000） -- 基于排序的实现. 也在 Cython 中。

Answer 2

如果 a 由整数组成 a 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，使用 collections.defaultdict 可能会更快。

编辑：参见 JF Sebastian 的解决方案 解决仅限整数值限制和大值问题的方法。

Answer 3

这种方法怎么样：

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