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