矢量化 np.arange 或等效

Vectorize np.arange or equivalent

我有一个很长的一维数组。我想创建一个数组,它是 np.arange() 应用于数组中每个值加上一些常量的结果。例如,如果常量 = 3 并且我的数组看起来像

[1,2,3,4,5]

我想得到

[[1,2,3]
 [2,3,4]
 [3,4,5]
 [4,5,6]
 [5,6,7]]

np.arange() 只接受标量作为参数。我试了一下 np.vectorize() 但没有成功。很明显,我可以用循环或列表来做到这一点,然后转换为数组,但我想知道是否有一个好的 numpy-only 解决方案。

您可以使用加法和广播:

>>> x = np.array([1,2,3,4,5])
>>> constant = 3
>>> x[:,None] + np.arange(constant)
array([[1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6],
       [5, 6, 7]])

也可以写成np.add.outer(x, np.arange(constant))