Векторизация 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.

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