numpy: cumulative multiplicity count

Viewed 482

I have a sorted array of ints which might have repetitions. I would like to count consecutive equal values, restarting from zero when a value is different from the previous one. This is the expected result implemented with a simple python loop:

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)

Output:

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]

How can I get the same result in an efficient way using numpy? The array is very long, but the repetitions are just a few (say no more than ten).

In my special case I also know that values start from zero and that the difference between consecutive values is either 0 or 1 (no gaps in values).

2 Answers
Related