高效地对2D NumPy数组的每一行应用不同的排列

Question

给定一个矩阵A，我想为A的每一行应用不同的随机洗牌；例如，

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

变成

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

当然我们可以遍历矩阵并对每一行进行随机重排；然而迭代速度较慢，我想知道是否有更高效的方法来完成这个任务。

- Tony

---

这里有另一个答案（https://dev59.com/TGEi5IYBdhLWcg3wueR0）。那里的评论也建议使用`apply_along_axis`。列的另一个答案在这里（https://stackoverflow.com/questions/26975807/efficient-way-to-shuffle-one-column-at-the-time-in-numpy-matrix）和（https://dev59.com/42Ij5IYBdhLWcg3wKyS0）以及（https://dev59.com/UJXfa4cB1Zd3GeqPgIIb）。 - Sheldore

还有一个这里也是关于列的。 - Sheldore

Answer 1

我从Divakar那里学到了这个巧妙的技巧，它涉及randn和argsort：

np.random.seed(0)

s = np.arange(16).reshape(4, 4)
np.take_along_axis(s, np.random.randn(*s.shape).argsort(axis=1), axis=1)

array([[ 1,  0,  3,  2],
       [ 4,  6,  5,  7],
       [11, 10,  8,  9],
       [14, 12, 13, 15]])

对于一个二维数组，这可以简化为

s[np.arange(len(s))[:,None], np.random.randn(*s.shape).argsort(axis=1)]

array([[ 1,  0,  3,  2],
       [ 4,  6,  5,  7],
       [11, 10,  8,  9],
       [14, 12, 13, 15]])

---

您还可以独立地对每一行应用np.random.permutation以返回一个新数组。

np.apply_along_axis(np.random.permutation, axis=1, arr=s)

array([[ 3,  1,  0,  2],
       [ 4,  6,  5,  7],
       [ 8,  9, 10, 11],
       [15, 14, 13, 12]])

性能 -

s = np.arange(10000 * 100).reshape(10000, 100) 

%timeit s[np.arange(len(s))[:,None], np.random.randn(*s.shape).argsort(axis=1)] 
%timeit np.apply_along_axis(np.random.permutation, 1, s)   

84.6 ms ± 857 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
842 ms ± 8.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我注意到这取决于你的数据维度，确保首先进行测试。

- cs95

---

谢谢！所以如果我有一个三维数组，如果我想要排列最后一个维度，那么我可以这样做 np.take_along_axis(s, np.random.randn(*s.shape).argsort(axis=2), axis=2)，对吗？ - Tony

@Tony 是的，我认为那应该可以。 - cs95

Answer 2

在编程方面，您可以使用numpy的apply_along_axis函数，如下所示：

np.apply_along_axis(np.random.shuffle, 1, matrix)

但对于3x3矩阵，似乎并不比迭代更有效，因为我得到的方法是

> %%timeit 
> np.apply_along_axis(np.random.shuffle, 1, test)
67 µs ± 1.8 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

当迭代进行时

> %%timeit
> for i in range(test.shape[0]):
>     np.random.shuffle(test[i])
20.3 µs ± 284 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

- vlizana

---

apply_along_axis 本质上只是在“其他”轴上进行迭代。没有速度保证。它使三维及更大尺寸的迭代更加美观；对于二维则无作用。 - hpaulj