Python 给定置换矩阵的numpy数组每行（或每列）的有效置换

让我们假设，我有两个给定的

ndarray

s，其中矩阵

mapping

包含矩阵

mask

行应如何排列的信息。我们可以假设，

映射

矩阵来自其他算法

import numpy as np
T, K, F = 2, 3, 5
mask = np.random.randint(4, size=(T, K, F))
mapping = np.asarray([
    [0, 1, 2],
    [0, 1, 2],
    [2, 0, 1],
    [0, 1, 2],
    [1, 0, 2]
])

要做到这一点，最直接的方法是应用

for

循环：

out = np.empty_like(mask)
for f in range(F):
    out[:, :, f] = mask[:, mapping[f, :], f]

这似乎相当有效，因此我研究并找到了以下解决方案：

out = mask[
    np.arange(T)[:, None, None],
    mapping.T[None, :, :],
    np.arange(F)[None, None, :]
]

相关问题的答案建议使用

ogrid

：

ogrid = np.ogrid[:T, :1, :F]
out = mask[
    ogrid[0],
    mapping.T[None, :, :],
    ogrid[2]
]

创建所有中间数组并正确广播它们似乎很不舒服。那么，执行所需的重新排序的最佳方法是什么

定时信息：

为了提供有意义的计时信息，我使用了一些更接近我的应用程序的形状。随机排列只是为了简单起见

T, K, F = 1000, 3, 257
mask = np.random.randint(4, size=(T, K, F))
mapping = np.stack([list(np.random.permutation(np.arange(3))) for _ in range(F)])

结果如下：

for loop:                 100 loops, best of 3: 8.4 ms per loop
three times broadcasting: 100 loops, best of 3: 8.37 ms per loop
ogrid:                    100 loops, best of 3: 8.33 ms per loop
swapaxis:                 100 loops, best of 3: 2.43 ms per loop
transpose:                100 loops, best of 3: 2.08 ms per loop

定义“最佳”是有争议的，但这里有一种方法-

---

另一种方法是转置

映射

，然后使用最后一个轴的范围数组，而不扩展到

2D

。

映射的每一行
最后一个轴（轴=-1）
，决定了
掩码的
第二个最后一个轴（轴=-2）
上元素的顺序。因此，我们需要
映射上的转置。在第一种方法中，我们通过后一种轴交换实现了这种转置行为。我可以从效率上保证这一点

因此，我们将有这样的实施-

mask[:,mapping.T, np.arange(F)]

我不太明白，为什么这会更快，但我添加了上面的计时，而且您的两个版本都明显更快。
mask[:,mapping.T, np.arange(F)]