python - 对 numpy 数组的非零元素进行排序并获取它们的索引

Question

我有一个 numpy 数组:

x =   numpy.array([0.1, 0, 2, 3, 0, -0.5])

我想得到一个数组 y，它包含已排序的 x 的非零元素和 idx，它是 x 的对应索引。

例如，对于上述示例，y 将为 [3, 2, 0.1, -0.5]，idx 将为 [3, 2, 0, 5]。我更喜欢一种可以扩展到二维数组而无需遍历 x 行的方法。

如果我有一个二维的例子

x = [[0.1, 0, 2, 3, 0, -0.5],
     [1, 0, 0, 0, 0, 2 ]] 

我想要一个

y =[[3, 2, 0.1, -0.5],[2,1]] and 
idx = [[3, 2, 0, 5], [5, 0]].

Answer 1

这是分别解决1D 和2D 情况的两种向量化方法 -

def sort_nonzeros1D(x):
    sidx = np.argsort(x)
    out_idx = sidx[np.in1d(sidx, np.flatnonzero(x!=0))][::-1]
    out_x = x[out_idx]
    return out_x, out_idx

def sort_nonzeros2D(x):
    x1 = np.where(x==0, np.nan, x)
    sidx = np.argsort(x1,1)[:,::-1]

    n = x.shape[1]
    extent_idx = (x==0).sum(1)
    valid_mask = extent_idx[:,None] <= np.arange(n)
    split_idx = (n-extent_idx[:-1]).cumsum()

    out_idx = np.split(sidx[valid_mask], split_idx)
    y = x[np.arange(x.shape[0])[:,None], sidx]
    out_x = np.split(y[valid_mask], split_idx)
    return out_x, out_idx

样本运行

1D 案例:

In [461]: x
Out[461]: array([ 0.1,  0. ,  2. ,  3. ,  0. , -0.5])

In [462]: sort_nonzeros1D(x)
Out[462]: (array([ 3. ,  2. ,  0.1, -0.5]), array([3, 2, 0, 5]))

2D 案例:

In [470]: x
Out[470]: 
array([[ 0.1,  0. ,  2. ,  3. ,  0. , -0.5],
       [ 1. ,  0. ,  0. ,  0. ,  0. ,  2. ],
       [ 7. ,  0. ,  2. ,  5. ,  1. ,  0. ]])

In [471]: sort_nonzeros2D(x)
Out[471]: 
([array([ 3. ,  2. ,  0.1, -0.5]),
  array([ 2.,  1.]),
  array([ 7.,  5.,  2.,  1.])],
 [array([3, 2, 0, 5]), array([5, 0]), array([0, 3, 2, 4])])