Python: moving all elements greater than 0 to left and right in numpy array

Question

If I have an numpy array like the one below, how can I right justify or left justify the elements tat are greater than zero

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

for example, If I wanted to right justify this array, it would look like:

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

Answer 1

One vectorized approach making use of 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

Basically the steps are :

• Make a mask of greater than zeros.

• Get a left or right justified mask where greater than elements are to be placed in a zeros initialized array. To get such a justified mask, we simply sort the mask from step-1 along each row, which throws the True ones in each row to the right. Thus, additionally we need flipping of each row for the left justified case.

• Finally, use the justified mask to assign into output array and the mask from step-1 to select from input array.

Sample runs -

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