如何使用二维索引数组按所需顺序放置 NumPy 数组的元素 [重复]

Question

我有一个格式为

的 NumPy 数组

a = [[[a1, a2, a3],
      [a4, a5, a6],
      [a7, a8, a9]],

     [[b1, b2, b3],
      [b4, b5, b6],
      [b7, b8, b9]]]

和另一个形式的数组

M = [[1, 0, 2],
     [1, 2, 0]]

有没有一种简单的方法可以使用第二个数组中的索引对第一个数组重新排序？所以最后第一个会变成：

a = [[[a4, a5, a6],
      [a1, a2, a3],
      [a7, a8, a9]],

     [[b4, b5, b6],
      [b7, b8, b9],
      [b1, b2, b3]]]

Answer 1

您可以使用列表推导式 -

b = [[a[i][j] for j in m] for i, m in enumerate(M)]

你也可以这样用：

b = [i[j] for i, j in zip(a, M)]

因为你的数组是 numpy 数组，但它会更慢。对于任何目的，列表理解都非常快。

如果您不关心一次加载整个列表，可以使用生成器：

b = ([a[i][j] for j in m] for i, m in enumerate(M))

但这只能迭代一次。

你也可以使用 numpy 方法：

b = a[np.arange(len(M))[:,None],M]

这将比对较大数组的列表推导更快。