在 numpy 中高效的每列矩阵索引 [英] efficient per column matrix indexing in numpy

问题描述

我有两个大小相同的矩阵，A、B.我想使用 B 的列来访问 A 的列，以每列为基础.例如，

I have two matrices of the same size, A, B. I want to use the columns of B to acsses the columns of A, on a per column basis. For example,

A = np.array([[1, 4, 7],
              [2, 5, 8],
              [3, 6, 9]])

和

B = np.array([[0, 0, 2],
              [1, 2, 1],
              [2, 1, 0]])

我想要类似的东西:

A[B] = [[1, 4, 9],  
        [2, 6, 8], 
        [3, 5, 7]]

即，我使用 B 的第 j 列作为 A 的第 j 列的索引.有没有什么有效的方法呢?谢谢！

I.e., I've used the j'th column of B as indices to the j'th column of A. Is there any effiecnt way of doing so? Thanks!

推荐答案

您可以使用 高级索引:

A[B, np.arange(A.shape[0])]

array([[1, 4, 9],
       [2, 6, 8],
       [3, 5, 7]])

或者使用 np.take_along_axis:

np.take_along_axis(A, B, axis=0)

array([[1, 4, 9],
       [2, 6, 8],
       [3, 5, 7]])

这篇关于在 numpy 中高效的每列矩阵索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！