python任意累积_numpy：累积多重计数

这是一种基于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