我有一个矩阵(准确地说是2d numpy ndarray):
A = np.array([[4, 0, 0],
[1, 2, 3],
[0, 0, 5]])
我想要滚动每一行 A
独立地,根据另一个数组中的滚动值:
r = np.array([2, 0, -1])
也就是说,我想这样做:
print np.array([np.roll(row, x) for row,x in zip(A, r)])
[[0 0 4]
[1 2 3]
[0 5 0]]
有没有办法有效地做到这一点?也许使用花哨的索引技巧?
当然你可以使用高级索引来做到这一点,无论它是否很好,可能取决于你的数组大小(如果你的行很大,可能不是):
rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]
# Use always a negative shift, so that column_indices are valid.
# (could also use module operation)
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:,np.newaxis]
result = A[rows, column_indices]
numpy.lib.stride_tricks.as_strided
stricks(简称双关语)再次!
说起 花哨的索引技巧,有的 臭名昭著 - np.lib.stride_tricks.as_strided
。想法/技巧将是从第一列开始到最后一列的切片部分,并在结束时连接。这确保了我们可以根据需要向前迈进,以便利用 np.lib.stride_tricks.as_strided
从而避免实际回滚的需要。这就是整个想法!
现在,就实际实施而言,我们将使用 scikit-image's view_as_windows
优雅地使用 np.lib.stride_tricks.as_strided
在引擎盖下。因此,最终的实施将是 -
from skimage.util.shape import view_as_windows as viewW
def strided_indexing_roll(a, r):
# Concatenate with sliced to cover all rolls
a_ext = np.concatenate((a,a[:,:-1]),axis=1)
# Get sliding windows; use advanced-indexing to select appropriate ones
n = a.shape[1]
return viewW(a_ext,(1,n))[np.arange(len(r)), (n-r)%n,0]
这是一个样本运行 -
In [327]: A = np.array([[4, 0, 0],
...: [1, 2, 3],
...: [0, 0, 5]])
In [328]: r = np.array([2, 0, -1])
In [329]: strided_indexing_roll(A, r)
Out[329]:
array([[0, 0, 4],
[1, 2, 3],
[0, 5, 0]])
标杆
# @seberg's solution
def advindexing_roll(A, r):
rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:,np.newaxis]
return A[rows, column_indices]
让我们对具有大量行和列的数组进行一些基准测试 -
In [324]: np.random.seed(0)
...: a = np.random.rand(10000,1000)
...: r = np.random.randint(-1000,1000,(10000))
# @seberg's solution
In [325]: %timeit advindexing_roll(a, r)
10 loops, best of 3: 71.3 ms per loop
# Solution from this post
In [326]: %timeit strided_indexing_roll(a, r)
10 loops, best of 3: 44 ms per loop