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

Question

我有一个大小为 (3, 5, 4) 的矩阵mat和一个大小为 (4,) 的向量vec ，其索引对应于矩阵的第一维（即介于 0 和 2 之间）。 我想提取一个大小为 (4, 5) 的数组，这可以通过mat[vec, :, [True] * len(vec)]来完成，但我想知道是否有使用 numpy 函数的更优雅的解决方案无需创建 boolean 值的新列表。

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

等效 - 它是否更优雅？

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

除非您想要 (4,5,4)，否则您必须为第一个和第三个维度提供相同大小的 arrays。 没有办法解决这个问题。