repeat indices with different repeat values in numpy

Question

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.

Answer 1

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]