矢量化np.arange或等效

Question

我有一个很长的1D阵列.我想创建一个数组,它是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解决方案.

Answer 1

你可以使用加法和广播:

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