CODEWITHSUNDEEP
python-3.xnumpyloopsmultidimensional-arrayvectorization
Akira
Akira

Reputation: 2870

Faster way to sort the array by finding the position of the first non-zero element in each row

I have a matrix in which I would like to sort this matrix based on the position of the first non-zero element in each row. For example, I would like to rearrange A in to B.

import numpy as np

A = np.array([[0, 0, 8, 7, 8],
              [0, 2, 2, 8, 9],
              [0, 0, 0, 6, 10],
              [5, 4, 4, 8, 10]])

B = np.array([[5, 4, 4, 8, 10],
              [0, 2, 2, 8, 9],
              [0, 0, 8, 7, 8],
              [0, 0, 0, 6, 10]])

To do so, I first find the position of the first non-zero element in each row: I first define the function position and then rearrange A based the permutation position(A).

def position(A):
    num_rows = A.shape[0]
    num_cols = A.shape[1]
    B = np.zeros(num_rows)
    # Access each row
    for i in range(num_rows):
      # Find the position of the first non-zero element
      for j in range(num_cols):
        if A[i, j] != 0:
          break
      B[i] = j
    return B

position(A)

array([2., 1., 3., 0.])

My matrix is of very large dimension and the procedure is repeated several times. My function position contains loops that I think could be replaced with more optimized code.

I would like to ask for a faster way to achieve my goal. Thank you so much!

Upvotes: 0

Views: 154

Answers (2)

Divakar
Divakar

Reputation: 221624

One way based on masking and argsort -

In [58]: m = A!=0

In [59]: A[m.argmax(1).argsort()]
Out[59]: 
array([[ 5,  4,  4,  8, 10],
       [ 0,  2,  2,  8,  9],
       [ 0,  0,  8,  7,  8],
       [ 0,  0,  0,  6, 10]])

If any row has all 0s, we need to modify the last step like so -

In [41]: A[0] = 0

In [42]: A[np.where(m.any(1),m.argmax(1),m.shape[1]-1).argsort()]
Out[42]: 
array([[ 5,  4,  4,  8, 10],
       [ 0,  2,  2,  8,  9],
       [ 0,  0,  0,  6, 10],
       [ 0,  0,  0,  0,  0]])

Upvotes: 1

Quang Hoang
Quang Hoang

Reputation: 150785

Not sure if this is fast, but sure is short:

B = np.array(sorted(A, key=lambda x: tuple(x!=0), reverse=True))

Output:

array([[ 5,  4,  4,  8, 10],
       [ 0,  2,  2,  8,  9],
       [ 0,  0,  8,  7,  8],
       [ 0,  0,  0,  6, 10]])

Upvotes: 1

Related Questions