python - 为 2D NumPy 数组的每一行高效应用不同的排列

Question

这个问题已经有答案了:

Randomly shuffle items in each row of numpy array (6 个回答)

已关闭 4 年前。

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

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

变成

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

当然我们可以循环遍历矩阵，让每一行随机打乱；但是迭代很慢，我想知道是否有更有效的方法来做到这一点。

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)

我注意到这取决于您数据的维度，请务必先进行测试。