Numpy 中的高效矩阵索引

Question

我有一个大小为 256 X 16 的矩阵A和一个大小为 8 X 8 的第二个矩阵B ，其中B的每个元素都是[0, 15]范围内的 integer 。 如果 A 只是一个行向量（即大小为 1X16），那么C = A[B]将 output 一个 8X8 矩阵，其中Cij = A[bij] 。 如果A是 256X16 使用 numpy （或其他任何东西），有什么有效的方法可以做同样的事情吗？ 我的意思是我想要一个大小为 256 X 8 X 8 的 output 矩阵C （即B正在索引A每一行）。 谢谢！

Answer 1

好的，我认为我解决了它：

C = A[:, B]