Ролл строки матрицы независимо друг от друга

Question

У меня есть матрица (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]] 

есть ли способ сделать это эффективно? Возможно, используя причудливые индексирующие трюки?

Answer 1

Конечно, вы можете сделать это с помощью передовых индексации, является ли это fastes способ, вероятно, вероятно, зависит от размера массива (если ваши строки являются большими, то не может быть) .:

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]