如何随机移动numpy数组的行

Question

我正在寻找一种更 Pythonic 的方式来随机移动 numpy 数组的行。这个想法是我有一个数据数组，我想将数组的每一行左移一个随机量。我的解决方案可行，但我觉得有点不符合 Python 风格：

def shift_rows(data, max_shift):
    """Left-shifts each row in `data` by a random amount up to `max_shift`."""
    return np.array([np.roll(row, -np.random.randint(0, max_shift)) for row in data])

然后进行测试：

data = np.array([np.arange(0, 5) for _ in range(10)])  # toy data to illustrate
shifted = shift_rows(data, max_shift=5)
shifted
# array([1, 2, 3, 4, 0],
#       [1, 2, 3, 4, 0],
#       [0, 1, 2, 3, 4],
#       ...
#       [4, 0, 1, 2, 3]])

这实际上更像是一个思想实验。任何人都可以想出一种更有效或更pythonic的方式来做到这一点吗？我想列表推导是 Python 的，但是如果我需要在一个巨大的数组上这样做，这效率高吗？

编辑：我将 Divakar 的出色回复标记为答案，但如果有人有任何其他想法，我仍然希望听到它。

Answer 1

一次性为所有行生成所有列索引，然后简单地使用integer-indexing 进行矢量化解决方案，就像这样 -

# Store shape of input array
m,n = data.shape

# Get random column start indices for each row in one go
col_start = np.random.randint(0, max_shift, data.shape[0])

# Get the rolled indices for every row again in a vectorized manner.
# We are extending col_start to 2D and then adding a range array to get 
# all column indices for every row by leveraging NumPy's braodcasting.
# Because of the additions, we might go off-limits. So, to simulate the 
# rolled over version, mod it.
idx = np.mod(col_start[:,None] + np.arange(n), n)

# Finall with integer indexing get the values off data array
shifted_out = data[np.arange(m)[:,None], idx]

分步运行-

1] 输入：

In [548]: data
Out[548]: 
array([[44, 23, 38, 32, 30],
       [69, 15, 32, 41, 63],
       [69, 41, 75, 50, 87],
       [23, 28, 38, 79, 91]])

In [549]: max_shift = 5

2] 建议的解决方案：

2A] 获取列开始：

In [550]: m,n = data.shape

In [551]: col_start = np.random.randint(0, max_shift, data.shape[0])

In [552]: col_start
Out[552]: array([1, 2, 3, 3])

2B] 获取所有索引：

In [553]: idx = np.mod(col_start[:,None] + np.arange(n), n)

In [554]: col_start[:,None]
Out[554]: 
array([[1],
       [2],
       [3],
       [3]])

In [555]: col_start[:,None] + np.arange(n)
Out[555]: 
array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7]])

In [556]: np.mod(col_start[:,None] + np.arange(n), n)
Out[556]: 
array([[1, 2, 3, 4, 0],
       [2, 3, 4, 0, 1],
       [3, 4, 0, 1, 2],
       [3, 4, 0, 1, 2]])

2C] 最后索引到数据：

In [557]: data[np.arange(m)[:,None], idx]
Out[557]: 
array([[23, 38, 32, 30, 44],
       [32, 41, 63, 69, 15],
       [50, 87, 69, 41, 75],
       [79, 91, 23, 28, 38]])

验证 -

1] 原始方法：

In [536]: data = np.random.randint(11,99,(4,5))
     ...: max_shift = 5
     ...: col_start = -np.random.randint(0, max_shift, data.shape[0])
     ...: for i,row in enumerate(data):
     ...:     print np.array([np.roll(row, col_start[i])])
     ...:     
[[83 93 17 53 61]]
[[55 88 84 94 89]]
[[59 63 29 72 85]]
[[57 95 13 21 14]]

2] 建议的方法重用col_start，以便我们可以进行价值验证：

In [537]: m,n = data.shape

In [538]: idx = np.mod(-col_start[:,None] + np.arange(n), n)

In [539]: data[np.arange(m)[:,None], idx]
Out[539]: 
array([[83, 93, 17, 53, 61],
       [55, 88, 84, 94, 89],
       [59, 63, 29, 72, 85],
       [57, 95, 13, 21, 14]])