如何利用二维行列索引数组重排二维NumPy数组

问题场景 #

给定原二维NumPy数组：

import numpy as np

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

以及对应形状的二维行索引数组和列索引数组：

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

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

需要生成目标数组：

result = np.array([[6, 4, 5],
                   [4, 6, 8],
                   [2, 8, 9]])

解决方案 #

直接使用NumPy的高级索引特性，一行代码即可完成：

result = a[rows, cols]

原理说明 #

当传入两个形状完全相同的二维索引数组时，NumPy会对结果数组的每个位置(i,j)，取原数组中a[rows[i,j], cols[i,j]]的值，完全匹配所需的重排规则。这种向量化操作比循环或map类方法效率高得多，尤其适合处理大规模数组。

验证结果 #

运行上述代码后，打印result会得到：

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

和目标结果完全一致。

内容的提问来源于stack exchange，提问作者shmulik90