python - 带有 numpy 的范围数组

Question

我有以下数组:

>>> x = numpy.array([2,4,2,3,1])
>>> x
array([2, 4, 2, 3, 1])

我想要一组这些值的范围。我可以这样创建它:

>>> numpy.hstack( (numpy.arange(v) for v in x) )
array([0, 1, 0, 1, 2, 3, 0, 1, 0, 1, 2, 0])

给定 x，有没有更快的方法来使用 numpy 生成它而无需使用 for 循环？

Answer 1

我想通了:

>>> x
array([2, 4, 2, 3, 1])
>>> ends = numpy.cumsum(x)
>>> ranges = numpy.arange(ends[-1])
>>> ranges = ranges - numpy.repeat(ends-x, x)
>>> ranges
array([0, 1, 0, 1, 2, 3, 0, 1, 0, 1, 2, 0])
>>>