可変距離上のNumPy行の効率的な回転

Question

2D M x N NumPy配列と回転距離のリストが与えられているので、すべてのM行をリスト内の距離にわたって回転させたいと思います。これは私が現在持っているものです。

import numpy as np 

M = 6 
N = 8 
dists = [2,0,2,1,4,2] # for example 
matrix = np.random.randint(0,2,(M,N)) 

for i in range(M): 
    matrix[i] = np.roll(matrix[i], -dists[i]) 

最後の2行は、実際に何百何千回ものを実行しますとcProfileによって測定されるように、それは私のパフォーマンスがボトルネックされる内部ループの一部です。例えばforループを回避し、これをより効率的に行うことは可能ですか？

Answer 1

range(0...N)の配列をdistsに追加した後でモジュラス演算を使ってローリングの振る舞いをシミュレートすると、同じ行で要素を選択してシャッフルする各行の列インデックスを得ることができます。 broadcastingの助けを借りて、すべての行にわたってこのプロセスをベクトル化できます。したがって、我々はそうのように実装しています - 私は `timeit`モジュールを搭載しただけで、これらのラインをテストしたときに、これは、ほぼ65％速いですが

M,N = matrix.shape # Store matrix shape 

# Get column indices for all elems for a rolled version with modulus operation 
col_idx = np.mod(np.arange(N) + dists[:,None],N) 

# Index into matrix with ranged row indices and col indices to get final o/p 
out = matrix[np.arange(M)[:,None],col_idx]