通过查找每行中第一个非零元素的位置来对数组进行排序的快速方法

Question

我有一个矩阵，我希望根据该矩阵在每一行中的第一个非零元素的位置对其进行排序。例如，我想将A重新排列为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]])

为此，我首先找到每行中第一个非零元素的位置：首先定义函数position，然后根据排列A重新排列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.])

我的矩阵尺寸很大，该过程重复了几次。我的函数position包含我认为可以用更优化的代码替换的循环。

我想寻求一种更快的方法来实现自己的目标。非常感谢！

Answer 1

不确定这是否很快，但肯定很短：

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

输出：

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

Answer 2

基于掩码和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]])

如果任何行全为0，我们需要像这样修改最后一步-

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]])