Numpy：使用 2d 数组沿 3d 数组的轴选择答案

【问题标题】：Numpy: select along axis of 3d array using 2d arrayNumpy：使用 2d 数组沿 3d 数组的轴选择
【发布时间】：2020-12-02 04:20:33
【问题描述】：

我已经为此苦苦挣扎了几个小时，无法完全理解它。设置是这样的：

    A.shape # (T,N,K)
    B.shape #   (L,K) L < N

2D B 数组的 K 列中的每一列都索引同一 K 行的 N 列之一。我可以通过

轻松抓取任何特定的 k 切片

    A[:,B[:,k],k].shape # (T,L)

但是，循环 K 并不理想，因为 A 是一个非常大的矩阵

我确信有人有一个非常简单的答案，但我很难过。

编辑：我还应该补充一点，我需要保留 A 矩阵的 3D 结构。我想出了如何获取单个值，但只能在 (TxLxK,) 数组中。

【问题讨论】：

标签： python numpy

【解决方案1】：

您可以使用np.take_along_axis

np.take_along_axis(A,B[None,...],axis=1)

例如，

A = np.linspace(1,24,24).reshape(3,4,2)
B = np.repeat([[0,1]],3,axis=0)

np.take_along_axis(A,B[None,...],axis=1)

结果是

array([[[ 1.,  4.],
        [ 1.,  4.],
        [ 1.,  4.]],

       [[ 9., 12.],
        [ 9., 12.],
        [ 9., 12.]],

       [[17., 20.],
        [17., 20.],
        [17., 20.]]])

【讨论】：