Search code examples
pythonnumpynumpy-ndarraydiagonal

Delete diagonals of zero elements


I'm trying to reshape an array from its original shape, to make the elements of each row descend along a diagonal:

np.random.seed(0) 
my_array = np.random.randint(1, 50, size=(5, 3))
array([[45, 48,  1],
       [ 4,  4, 40],
       [10, 20, 22],
       [37, 24,  7],
       [25, 25, 13]])

I would like the result to look like this:

my_array_2 = np.array([[45,  0,  0],
                       [ 4, 48,  0],
                       [10,  4,  1],
                       [37, 20, 40],
                       [25, 24, 22],
                       [ 0, 25,  7],
                       [ 0,  0, 13]])

This is the closest solution I've been able to get:

my_diag = []
for i in range(len(my_array)):
    my_diag_ = np.diag(my_array[i], k=0)
    my_diag.append(my_diag_)
my_array1 = np.vstack(my_diag)
array([[45,  0,  0],
       [ 0, 48,  0],
       [ 0,  0,  1],
       [ 4,  0,  0],
       [ 0,  4,  0],
       [ 0,  0, 40],
       [10,  0,  0],
       [ 0, 20,  0],
       [ 0,  0, 22],
       [37,  0,  0],
       [ 0, 24,  0],
       [ 0,  0,  7],
       [25,  0,  0],
       [ 0, 25,  0],
       [ 0,  0, 13]])

From here I think it might be possible to remove all zero diagonals, but I'm not sure how to do that.


Solution

  • In [134]: arr = np.array([[45, 48,  1],
         ...:        [ 4,  4, 40],
         ...:        [10, 20, 22],
         ...:        [37, 24,  7],
         ...:        [25, 25, 13]])
    In [135]: res= np.zeros((arr.shape[0]+arr.shape[1]-1, arr.shape[1]), arr.dtype)
    

    Taking a hint from how np.diag indexes a diagonal, iterate on the rows of arr:

    In [136]: for i in range(arr.shape[0]):
         ...:     n = i*arr.shape[1]
         ...:     m = arr.shape[1]
         ...:     res.flat[n:n+m**2:m+1] = arr[i,:]
         ...: 
    In [137]: res
    Out[137]: 
    array([[45,  0,  0],
           [ 4, 48,  0],
           [10,  4,  1],
           [37, 20, 40],
           [25, 24, 22],
           [ 0, 25,  7],
           [ 0,  0, 13]])