使用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来生成这个，而不需要使用循环呢？

Answer 1

这个真的更快吗？

我有类似的需求，

concatenate([range(l, r) for l, r in array((left, right)).T])

的速度是

range(end[-1]) + repeat(left + end, right-left)

的两倍。

（这里的end = cumsum(right - left)和你的用法一样。）

（根据我短暂的经验，repeat的速度非常慢——至少在python 3.6中是这样的。）

Answer 2

我搞明白了：

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