Roll rows of a matrix independently

numpy, performance, python

Solution

You can do it using advanced indexing. Whether or not it is the fastest way likely depends on the array size. For instance, for large rows, this may be slower than other methods.

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Always use a negative shift, so that column_indices are valid.
# Alternative: r %= A.shape[1]
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:, np.newaxis]

result = A[rows, column_indices]

Problem

I have a matrix (2d numpy ndarray, to be precise): ``` A = np.array([[4, 0, 0], [1, 2, 3], [0, 0, 5]]) ``` And I want to roll each row of `A` independently, according to roll values in another array: ``` r = np.array([2, 0, -1]) ``` That is, I want to do this: ``` print np.array([np.roll(row, x) for row,x in zip(A, r)]) [[0 0 4] [1 2 3] [0 5 0]] ``` Is there a way to do this efficiently? Perhaps using fancy indexing tricks?

Original source