python - Numpy - slicing 2d row or column vector from array

Question

I'm trying to find a neat little trick for slicing a row/column from a 2d array and obtaining an array of (col_size x 1) or (1 x row_size).

Is there an easier way than to use numpy.reshape() after every slicing?

Cheers, Stephan

Answer 1

You can slice and insert a new axis in one single operation. For example, here's a 2D array:

>>> a = np.arange(1, 7).reshape(2, 3)
>>> a
array([[1, 2, 3],
       [4, 5, 6]])

To slice out a single column (returning array of shape (2, 1)), slice with None as the third dimension:

>>> a[:, 1, None]
array([[2],
       [5]])

To slice out a single row (returning array of shape (1, 3)), slice with None as the second dimension:

>>> a[0, None, :]
array([[1, 2, 3]])