numpy中每列矩阵索引的高效

Question

我有两个大小相同的矩阵 A、B。我想使用 B 的列来逐列访问 A 的列。例如，

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 列的索引。 有什么有效的方法吗？ 谢谢！

Answer 1

你可以使用advanced indexing:

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