生成连接范围的一维 NumPy 数组答案

【问题标题】：Generate 1D NumPy array of concatenated ranges生成连接范围的一维 NumPy 数组
【发布时间】：2017-05-09 16:06:55
【问题描述】：

我要生成如下数组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, ...]

有没有更好的方法？

【问题讨论】：

错字 - np.range 应该是 np.arange，我相信。
np.r_[:4, :5, :3, :6] 是“手动”生成这样一个数组的一种很好的紧凑方式。如果从 nv 数组开始，它不会提供任何优势。

标签： python arrays numpy python-3.5

【解决方案1】：

这是使用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

【讨论】：

哦，这么快，现在我必须为你的答案投票（尽管我不完全理解 magic id_arr[shifts[:-1]] = -nv[:-1]+1 行）
@MSeifert 是的，不是直截了当。稍后可能会解释。

【解决方案2】：

为什么要分两步？

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

【讨论】：

我相信第一步只是创建一个示例输入。
是的，这只是一个示例输入。
知道了。好吧，我想我会把它留在这里，以防其他人没有意识到问题是从现有的长度列表构建范围数组。

【解决方案3】：

您可以使用普通的 python ranges 并将结果转换为数组，而不是使用 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)

这不需要使用难以理解的 numpy 切片和构造。展示一个示例案例：

import random

>>> nv = [random.randint(1, 10) for _ in range(5)]
>>> print(nv)
[4, 2, 10, 5, 3]

>>> print(your_func(nv))
[ 1  2  3  4  1  2  1  2  3  4  5  6  7  8  9 10  1  2  3  4  5  1  2  3]

【讨论】：