Search code examples
pythonnumpymatrixmatrix-transform

matrix: move n-th row by n position efficiently


I have a numpy 2d array and I need to transform it in a way that the first row remains the same, the second row moves by one position to right, (it can wrap around or just have zero padded to the front). Third row shifts 3 positions to the right, etc. I can do this through a "for loop" but that is not very efficient. I am guessing there should be a filtering matrix that multipled by the original one will have the same effect, or maybe a numpy trick that will help me doing this? Thanks! I have looked into numpy.roll() but I don't think it can work on each row separately.

import numpy as np
p = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
'''
p = [ 1   2   3   4
      5   6   7   8
      9   10  11  12
      13  14  15  16]
desired output:
p'= [ 1   2   3   4
      0   5   6   7
      0   0   9   10
      0   0   0   13]
'''

Solution

  • Here is a simple method based on zero-padding and reshaping. It is fast because it avoids advanced indexing and other overheads.

    def pp(p):
        m,n = p.shape
        aux = np.zeros((m,n+m-1),p.dtype)
        np.copyto(aux[:,:n],p)
        return aux.ravel()[:-m].reshape(m,n+m-2)[:,:n].copy()