Shift rows of a numpy array independently

Question

This is an extension of the question posed here (quoted below)

I have a matrix (2d numpy ndarray, to be precise):

A = np.array([[4, 0,         


        

Answer 1

Inspired by Roll rows of a matrix independently's solution, here's a vectorized one based on 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
    p = np.full((a.shape[0],a.shape[1]-1),np.nan)
    a_ext = np.concatenate((p,a,p),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)), -r + (n-1),0]

Sample run -

In [76]: a
Out[76]: 
array([[4, 0, 0],
       [1, 2, 3],
       [0, 0, 5]])

In [77]: r
Out[77]: array([ 2,  0, -1])

In [78]: strided_indexing_roll(a, r)
Out[78]: 
array([[nan, nan,  4.],
       [ 1.,  2.,  3.],
       [ 0.,  5., nan]])