Question

我有一个排序的整数数组，可能有重复。我想计算连续的相等值，当值与前一个值不同时，从零重新开始。这是使用简单的python循环实现的预期结果：

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以有效的方式获得相同的结果？阵列很长，但重复只有几个（比如说不超过十个）。

在我的特殊情况下，我也知道值从零开始，连续值之间的差异为0或1（值没有间隙）。

Answer 1

这是一种基于cumsum的矢量化方法 -

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

运行时测试 -

{{1}}

Related post

Answer 2

我会在这些问题上使用numba

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

根本不重写代码，它比Divakar的矢量化解决方案快50％。

如果导致更短且可能更容易理解的代码，那么矢量化很有用，但如果你强行必须对代码进行矢量化，这对于一个相当经验的程序员来说也是一个问题，那么numba是可行的方法。

numpy：累积多重性计数

2 个答案: