2011-09-24 131 views
5

鉴于相同的长度a和b两个无序阵列:Python的组由阵列的,并总结数组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] 

可替换地,随机地和双向蟒蛇:

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

更新:Approach3,由巴勃罗。时间:1.0265750885

def approach_Pablo(a,b):  

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

更新:方法4,由Unutbu。时间:0.184849023819 [WINNER至今,但作为唯一的整数]

def unique_Unutbu(a,b): 

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

也许有人发现一个更聪明的办法解决这个问题比我:)

+0

什么是无序数组? –

+0

我的意思是你不能假定列表a已排序。 – Helga

回答

5

如果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具有很大的长度,但不是很大的值。对于长度较小(大或小的值)的数组a,使用collections.defaultdict可能会更快。

编辑:请参阅J.F. Sebastian's solution以解决整数值限制和大值问题。

+0

[测量](https://gist.github.com/da57326584a3811652aa#file_measurements.org)显示'np.bincount()即使对[基于Cython的解决方案]也表现良好(https://gist.github.com/ da57326584a3811652aa#file_pdf.pyx)。 – jfs

2

如何对这种做法:

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

你可以使用defaultdict(int),它更干净。 –

+0

谢谢!我不知道。更新回答:) – Pablo

+0

我喜欢这个方法,很漂亮。不幸的是,它似乎比'方法1'更慢,尤其是对于长阵列... – Helga

6

这里的类似的方法来@unutbu's one

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倍。

我用另外三种方法制作了performance comparison。领导者是:用于整数阵列--Cython中的hash-based implementation;对于double数组(输入大小为10000) - sort-based impl.也在Cython中。