生成连续范围的1D NumPy数组

时间:2016-12-24 03:38:32

标签: python arrays numpy python-3.5

我想生成以下数组a

nv = np.random.randint(3, 10+1, size=(1000000,))
a = np.concatenate([np.arange(1,i+1) for i in nv])

因此,输出类似于 -

[0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 0, ...]

有没有更好的方法呢?

3 个答案:

答案 0 :(得分:2)

这是使用cumulative summation -

的矢量化方法
def ranges(nv, start = 1):
    shifts = nv.cumsum()
    id_arr = np.ones(shifts[-1], dtype=int)
    id_arr[shifts[:-1]] = -nv[:-1]+1
    id_arr[0] = start # Skip if we know the start of ranges is 1 already
    return id_arr.cumsum()

样品运行 -

In [23]: nv
Out[23]: array([3, 2, 5, 7])

In [24]: ranges(nv, start=0)
Out[24]: array([0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6])

In [25]: ranges(nv, start=1)
Out[25]: array([1, 2, 3, 1, 2, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7])

运行时测试 -

In [62]: nv = np.random.randint(3, 10+1, size=(100000,))

In [63]: %timeit your_func(nv) # @MSeifert's solution
10 loops, best of 3: 129 ms per loop

In [64]: %timeit ranges(nv)
100 loops, best of 3: 5.54 ms per loop

答案 1 :(得分:1)

您可以使用普通的python Sneaky并将结果转换为数组,而不是使用numpy方法执行此操作:

range

这并不需要利用难以理解的numpy切片和构造。要显示示例案例:

from itertools import chain
import numpy as np

def your_func(nv):
    ranges = (range(1, i+1) for i in nv)
    flattened = list(chain.from_iterable(ranges))
    return np.array(flattened)

答案 2 :(得分:0)

为什么要两步?

a = np.concatenate([np.arange(0,np.random.randint(3,11)) for i in range(1000000)])