在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]