Search code examples
pythonnumpycupy

Roll first column by 1, second column by 2, etc

I have an array in numpy. I want to roll the first column by 1, second column by 2, etc.

Here is an example.

>>> x = np.reshape(np.arange(15), (5, 3))
>>> x
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

What I want to do:

>>> y = roll(x)
>>> y
array([[12, 10,  8],
       [ 0, 13, 11],
       [ 3,  1, 14],
       [ 6,  4,  2],
       [ 9,  7,  5]])

What is the best way to do it?

The real array will be very big. I'm using cupy, the GPU version of numpy. I will prefer solution fastest on GPU, but of course, any idea is welcomed.

Solution

  • You could use advanced indexing:

    import numpy as np

x = np.reshape(np.arange(15), (5, 3))

h, w = x.shape

rows, cols = np.arange(h), np.arange(w)
offsets = cols + 1
shifted = np.subtract.outer(rows, offsets) % h

y = x[shifted, cols]

    y:

    array([[12, 10,  8],
       [ 0, 13, 11],
       [ 3,  1, 14],
       [ 6,  4,  2],
       [ 9,  7,  5]])