Update selected columns in 2D numpy array with 1D array

Question

Given the following arrays:

from numpy import * 
b = ones((5,5))
a = arange(4)

How do I get the following array with minimum amount of code? Basically update parts of array b with array a :

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

In matlab I can use one line to achieve this:

b = ones(5,5);
a = [0,1,2,3];
b(1:4,2:4) = repmat(a',[1,3]) 

Answer 1

You can write:

b[0:4, 1:4] = a[:, None]

Which makes b equal to:

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

b[0:4, 1:4] selects the appropriate slice of b (recall that Python uses zero-based indexing).

To complete the assignment of the vector a , it is necessary to add an extra axis of length 1 using a[:, None] . This is because the slice of b has shape (4, 3) and we need a to have shape (4, 1) so that the axes line up correctly to allow broadcasting.

Answer 2

Initialize output array and set a , just like we did for MATLAB -

b = np.ones((5,5))
a = np.array([0,1,2,3])

Now, let's use the automatic broadcasting supported by NumPy to replace the explicit replication done by repmat in MATLAB, for which we need to make a a 2D array by "pushing" the 1D elements along the first axis and introducing a singleton dimension as the second axis with np.newaxis as a[:,np.newaxis] . Please note the general term for dimension in NumPy is axis. A shorthand for np.newaxis is None , thus we need to use a[:,None] and use this for assigning elements into b .

Thus, the final step would be considering we have 0-based indexing in Python, we would have -

b[0:4,1:4] = a[:,None]