0

I have a numpy vector of counts such as this

[2, 4, 3]

and I am trying to create an array of its indexes

[0, 1, 2]

with each index repeated the number times specified by the corresponding count, ie, the final output should be

[0, 0, 1, 1, 1, 1, 2, 2, 2]

It's easy but inefficient to do this using a for loop. Is there a fast vectorized solution?

user3124206
  • 375
  • 1
  • 7
  • 16
  • 1
    Use [`np.repeat`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.repeat.html) (e.g. `np.repeat([0, 1, 2], [2, 4, 3])`). In plain Python, `[i for i, r in zip([0, 1, 2], [2, 4, 3]) for _ in range(r)]`. – jdehesa Nov 01 '19 at 18:23
  • 1
    Though it may not be possible for your use case, I think the best solution would be to try to not create that expanded list in the first place, and modify whatever code expects the list to accept the counts. – jedwards Nov 01 '19 at 18:24

0 Answers0