python 二维数组切片,numpy的 - 从切片二维数组的行或列向量

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

解决方案

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