CODEWITHSUNDEEP
pythonnumpyindexingvectorizationnumpy-ndarray
dhokas
dhokas

Reputation: 1789

repeat indices with different repeat values in numpy

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.

Upvotes: 3

Views: 821

Answers (1)

tel
tel

Reputation: 13999

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]

Upvotes: 4

Related Questions