sams-studio's questions

假设我有一个排列数组perm，如下所示：

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

如果我想将其应用于一个轴，我可以写类似如下的内容：

v = np.arange(9).reshape(3, 3)
print(v[perm])

输出：

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

       [[3, 4, 5],
        [6, 7, 8],
        [0, 1, 2]],

       [[0, 1, 2],
        [6, 7, 8],
        [3, 4, 5]],

       [[6, 7, 8],
        [3, 4, 5],
        [0, 1, 2]]])

现在我想同时将其应用于两个轴。我发现我可以通过以下方式实现：

np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

但我发现它效率很低，因为它必须创建一个网格，而且还需要一个 for 循环。我在这个例子中创建了一个小数组，但实际上我有很多更大的数组和很多排列，所以我真的很想有一个像单轴版本一样快速和简单的版本。