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