2017-07-26 78 views
3

我有一個可能有重複的整數的整數數組。我想計算連續的相等值,當值與前一值不同時,從零開始重新計數。這是預期的結果實現了簡單的Python循環:numpy:累計重複次數

import numpy as np 

def count_multiplicities(a): 
    r = np.zeros(a.shape, dtype=a.dtype) 
    for i in range(1, len(a)): 
     if a[i] == a[i-1]: 
      r[i] = r[i-1]+1 
     else: 
      r[i] = 0 
    return r 

a = (np.random.rand(20)*5).astype(dtype=int) 
a.sort() 

print "given sorted array: ", a 
print "multiplicity count: ", count_multiplicities(a) 

輸出:

given sorted array: [0 0 0 0 0 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4] 
multiplicity count: [0 1 2 3 4 0 1 2 0 1 2 3 0 1 2 3 0 1 2 3] 

我怎樣才能得到相同的結果,用numpy的一種有效的方法是什麼?陣列非常長,但重複只是少數(不超過10次)。

在我的特殊情況下,我也知道值從零開始,並且連續值之間的差值爲0或1(值中沒有差距)。

+0

進一步的想法:也許一個可能的解決方案可以通過使用與給定對角線和上對角線的良好選擇矩陣的乘法來實現。 –

回答

3

這裏有一個cumsum基於矢量的方法 -

def count_multiplicities_cumsum_vectorized(a):  
    out = np.ones(a.size,dtype=int) 
    idx = np.flatnonzero(a[1:] != a[:-1])+1 
    out[idx[0]] = -idx[0] + 1 
    out[0] = 0 
    out[idx[1:]] = idx[:-1] - idx[1:] + 1 
    np.cumsum(out, out=out) 
    return out 

採樣運行 -

In [58]: a 
Out[58]: array([0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4]) 

In [59]: count_multiplicities(a) # Original approach 
Out[59]: array([0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2]) 

In [60]: count_multiplicities_cumsum_vectorized(a) 
Out[60]: array([0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2]) 

運行測試 -

In [66]: a = (np.random.rand(200000)*1000).astype(dtype=int) 
    ...: a.sort() 
    ...: 

In [67]: a 
Out[67]: array([ 0, 0, 0, ..., 999, 999, 999]) 

In [68]: %timeit count_multiplicities(a) 
10 loops, best of 3: 87.2 ms per loop 

In [69]: %timeit count_multiplicities_cumsum_vectorized(a) 
1000 loops, best of 3: 739 µs per loop 

Related post

+0

謝謝!我所缺少的是'flatnonzero' –

1

我想這樣的問題

import numba 
nb_count_multiplicities = numba.njit("int32[:](int32[:])")(count_multiplicities) 
X=nb_count_multiplicities(a) 

使用numba沒有在所有重寫你的代碼,它比Divakar的矢量化的解決方案快大約50%。

如果矢量化的結果是代碼更短,也更容易理解,那麼矢量化會有很多用處,但是如果你強制必須矢量化代碼,這對於一個非常有經驗的程序員來說可能是一個問題,那麼Numba就是要走的路。

+0

很高興知道關於numba! –