повторять индексы с разными значениями повторения в numpy

Question

Я ищу эффективный способ сделать следующее с Numpy:

Учитывая массив counts натуральных чисел, содержащий, например:

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

Я хотел бы создать другой массив, содержащий индексы первого, где индекс i повторяется counts[i] раз:

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

Моя проблема в том, что этомассив потенциально очень большой, и я ищу векторный (или быстрый) способ сделать это.

Answer 1

Вы можете сделать это с помощью 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)

Выход:

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