Получить индексы ненулевых значений в массиве numpy с учетом повторяющихся индексов со значением> 1

Question

Учитывая

>>> a = np.array((0, 1, 0, 0, 0, 2, 0, 1))

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

>>> np.nonzero(a)  # or ~equivalently np.argwhere(a)
(array([1, 5, 7]),)

Однако на самом деле мне нужно (array([1, 5, 5, 7]),), потому что значение в a[5] равно 2 - т.е. значение a в индексе - это то, сколько раз этот индекс должен появиться в выводе.

Мое текущее мнение, поскольку максимальное значение не будет очень большим, что делать:

# pseudo
indices = np.nonzero(a)
while extra_indices:
    a -= 1
    extra_indices = np.nonzero(a)
    indices = np.concat(indices, extra_indices)

Есть способ лучше? (Хотя максимальное число не будет большим, массивы могут быть довольно большими, поэтому я бы предпочел убрать шаг итерации)

Answer 1

Вы можете просто repeat:

idx = np.nonzero(a)

out = np.repeat(idx,a[idx])
# array([1, 5, 5, 7])