python - numpy 的 bincount 函数的反函数

Question

给定一个整数数组 c，我如何将它转换成一个整数数组 inds 使得 np.all(np.bincount(inds ) == c) 是真的吗？

例如:

>>> c = np.array([1,3,2,2])
>>> inverse_bincount(c)  # <-- what I need

array([0,1,1,1,2,2,3,3])

上下文:我正在尝试跟踪多组数据的位置，同时对所有数据集同时执行计算。我将所有数据连接在一起进行批处理，但我需要一个索引数组来提取结果。

当前的解决方法:

def inverse_bincount(c):
  return np.array(list(chain.from_iterable([i]*n for i,n in enumerate(c))))

Answer 1

使用 numpy.repeat :

np.repeat(np.arange(c.size), c)