6

I want to shift each row by its row-number and be relative to my desired output shape. An example:

array([[0, 1, 2],        array([[0, 1, 2],                array([[0, 1, 2, 0, 0],
       [1, 2, 3],    ->            [1, 2, 3],      ->            [0, 1, 2, 3, 0],
       [2, 3, 4]])                    [2, 3, 4]])                [0, 0, 2, 3, 4])

The array to furthest left is my input and the array to the furthest right is my desired output. This can be generalized to bigger arrays, for example a 10x10 array.

Is there a nice way of doing this?


What I have is:

A = np.array([[1, 2, 3],
              [2, 3, 4],
              [3, 4, 5]], dtype=np.float32)

out = np.zeros((A.shape[0], A.shape[1]*2-1))

out[[np.r_[:3]], [np.r_[:3] + np.r_[:3][:,None]]] = A
BigBen
  • 46,229
  • 7
  • 24
  • 40
Kevin
  • 3,096
  • 2
  • 8
  • 37

2 Answers2

2

Here is a solution for square matrices of size n.

np.concatenate((A,np.zeros((n,n))),axis=1).flatten()[0:-n].reshape([n,2*n-1])

Javier A
  • 539
  • 2
  • 12
  • 1
    Thanks! This seems faster than using advanced indexing. I also noticed that your solution would be a bit faster by using ```numpy.ravel``` instead of ```np.flatten```. – Kevin Apr 30 '21 at 13:58
2

Here is a fast way:

def f(a):
    r, _ = a.shape
    return np.c_[a, np.zeros((r, r), dtype=a.dtype)].ravel()[:-r].reshape(r, -1)

Example:

a = np.arange(8).reshape(4, -1)
>>> f(a)
array([[0, 1, 0, 0, 0],
       [0, 2, 3, 0, 0],
       [0, 0, 4, 5, 0],
       [0, 0, 0, 6, 7]])

Timing

a = np.random.randint(100, size=(1000, 1000))
%timeit f(a)
# 2.08 ms ± 14.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Pierre D
  • 24,012
  • 7
  • 60
  • 96