Python:在numpy数组中左右移动所有大于0的元素

Question

如果我有一个像下面这样的 numpy 数组，我如何右对齐或左对齐大于零的元素

[[ 0.  5.  0.  2.]
 [ 0.  0.  3.  2.]
 [ 0.  0.  0.  0.]
 [ 2.  0.  0.  1.]]

例如，如果我想右对齐这个数组，它看起来像:

[[ 5.  2.  0.  0.]
 [ 3.  2.  0.  0.]
 [ 0.  0.  0.  0.]
 [ 2.  1.  0.  0.]]

Answer 1

一种利用 masks 的矢量化方法-

def justify_rows(a, side='left'):
    mask = a>0
    justified_mask = np.sort(mask,1)
    if side=='left':
        justified_mask = justified_mask[:,::-1]
    out = np.zeros_like(a) 
    out[justified_mask] = a[mask]
    return out

基本上步骤是:

• 制作大于零的掩码。

• 获取左对齐或右对齐掩码，其中大于的元素将放置在零初始化数组中。为了得到这样一个合理的掩码，我们只需沿着每一行对步骤 1 中的掩码进行排序，这会将每行中的 True 扔到右边。因此，对于左对齐的情况，我们还需要翻转每一行。

• 最后，使用对齐的掩码分配给输出数组，并使用步骤 1 中的掩码从输入数组中进行选择。

样本运行-

In [105]: a
Out[105]: 
array([[ 0.,  5.,  0.,  2.],
       [ 0.,  0.,  3.,  2.],
       [ 0.,  0.,  0.,  0.],
       [ 2.,  0.,  0.,  1.]])

In [106]: justify_rows(a, side='left')
Out[106]: 
array([[ 5.,  2.,  0.,  0.],
       [ 3.,  2.,  0.,  0.],
       [ 0.,  0.,  0.,  0.],
       [ 2.,  1.,  0.,  0.]])

In [107]: justify_rows(a, side='right')
Out[107]: 
array([[ 0.,  0.,  5.,  2.],
       [ 0.,  0.,  3.,  2.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  2.,  1.]])