如何将 numpy 向量转换为矩阵，其中矩阵中的每一列都包含初始向量中各个元素周围的范围？

Question

Say I have a numpy vector array:

array([1, 2, 3])

and I want to convert this vector into a matrix where each column takes a range of +/- 2 around the respective element in the initial vector, such that my output matrix is:

array([[-1,  0,  1],
       [ 0,  1,  2],
       [ 1,  2,  3],
       [ 2,  3,  4],
       [ 3,  4,  5]])

what is the best (preferably vectorized) way to do this?

Answer 1

You can do it with the following one-liner:

result = a + np.array([-1, 0, 1])[:, np.newaxis]

(I think, a more elegant solution).

The result is:

array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4]])

Edit

If the range depends on some parameter, say rng , you can do it as:

rng = 2   # From x-2 to x+2
result = a + np.arange(-rng, rng + 1)[:, np.newaxis]

getting:

array([[-1,  0,  1],
       [ 0,  1,  2],
       [ 1,  2,  3],
       [ 2,  3,  4],
       [ 3,  4,  5]])