快速纽姆

Question

我有一个2dnumpy数组，我想以增量方式滚动每一行。我在for循环中使用np.roll来执行此操作。但是由于我调用这个函数上千次，我的代码非常慢。你能帮我讲讲如何使它更快吗。在

我的输入看起来像

array([[4,1],
       [0,2]])

我的输出看起来像

^{pr2}$

这里第零行[4,1]被移动了0，第一行[0,2]被移动了1。同样，第二行将被移位2，以此类推。在

编辑

temp = np.zeros([dd,dd])
for i in range(min(t + 1, dd)):
    temp[i,:] = np.roll(y[i,:], i, axis=0)

Answer 1

这里有一个矢量化的解决方案-

m,n = a.shape
idx = np.mod((n-1)*np.arange(m)[:,None] + np.arange(n), n)
out = a[np.arange(m)[:,None], idx]

样本输入，输出-

^{pr2}$

因为，您已经提到要多次调用这样一个滚动例程，请创建一次索引数组idx，稍后再使用它。在

进一步改进

对于重复使用，最好先创建完整的线性索引，然后使用np.take来提取滚动元素，如下-

full_idx = idx + n*np.arange(m)[:,None]
out = np.take(a,full_idx)

让我们看看有什么改进-

In [330]: a = np.random.randint(11,99,(600,600))

In [331]: m,n = a.shape
     ...: idx = np.mod((n-1)*np.arange(m)[:,None] + np.arange(n), n)
     ...: 

In [332]: full_idx = idx + n*np.arange(m)[:,None]

In [333]: %timeit a[np.arange(m)[:,None], idx] # Approach #1
1000 loops, best of 3: 1.42 ms per loop

In [334]: %timeit np.take(a,full_idx)          # Improvement
1000 loops, best of 3: 486 µs per loop

关于3x的改进！在

Answer 2

一个棘手但快速的解决方案：

p=5
a=randint(0,p,(p,p))

aa=hstack((a,a))
m,n=aa.strides

b=np.lib.stride_tricks.as_strided(aa,a.shape,(m+n,n)) 
c=np.lib.stride_tricks.as_strided(aa.ravel()[p:],a.shape,(m-n,n)) 
  ## 

[[2 1 4 2 4]
 [0 4 2 0 3]
 [1 3 3 4 4]
 [1 0 3 2 4]
 [3 3 2 1 3]]

[[2 1 4 2 4]
 [4 2 0 3 0]
 [3 4 4 1 3]
 [2 4 1 0 3]
 [3 3 3 2 1]]

[[2 1 4 2 4]
 [3 0 4 2 0]
 [4 4 1 3 3]
 [3 2 4 1 0]
 [3 2 1 3 3]]