repeat indices with different repeat values in numpy

user_2633673 0
Viewed 440

I'm looking for an efficient way to do the following with Numpy:

Given a array counts of positive integers containing for instance:

[3, 1, 0, 6, 3, 2]

I would like to generate another array containing the indices of the first one, where the index i is repeated counts[i] times:

[0 0 0 1 3 3 3 3 3 3 4 4 4 5 5]

My problem is that this array is potentially very large and I'm looking for a vectorial (or fast) way to do this.

vectorization numpy-ndarray indexing numpy python
1 Answers

You can do it with numpy.repeat:

import numpy as np

arr = np.array([3, 1, 0, 6, 3, 2])
repix = np.repeat(np.arange(arr.size), arr)
print(repix)

Output:

[0 0 0 1 3 3 3 3 3 3 4 4 4 5 5]
user_4254580
Related