Numpy function 通过索引向量从 ndarray 中提取值

Question

I have a matrix mat of size (3, 5, 4) and a vector vec of size (4,) with indices corresponding to the first dimension of the matrix (ie between 0 and 2). I would like to extract an array of size (4, 5), which can be done via mat[vec, :, [True] * len(vec)] , but I was wondering if there is a more elegant solution using numpy functions without the need to create a new list of boolean values.

Answer 1

In [15]: mat = np.arange(3 * 5 * 4).reshape(3, 5, 4)
In [16]: idx = np.array([0, 2, 1, 1])

In [18]: mat[idx, :, [True] * len(idx)]
Out[18]: 
array([[ 0,  4,  8, 12, 16],
       [41, 45, 49, 53, 57],
       [22, 26, 30, 34, 38],
       [23, 27, 31, 35, 39]])

equivalent - whether it's more elegant?

In [19]: mat[idx, :, np.arange(4)]
Out[19]: 
array([[ 0,  4,  8, 12, 16],
       [41, 45, 49, 53, 57],
       [22, 26, 30, 34, 38],
       [23, 27, 31, 35, 39]])

Unless you want a (4,5,4), you will have to provide equal size arrays for the 1st and 3rd dimensions. There's no way around that.