Python: Justifying NumPy array

Here's a vectorized approach inspired by this other post and generalized to cover non-zeros for all four directions -

def justify(a, invalid_val=0, axis=1, side='left'):        """    Justifies a 2D array    Parameters    ----------    A : ndarray        Input array to be justified    axis : int        Axis along which justification is to be made    side : str        Direction of justification. It could be 'left', 'right', 'up', 'down'        It should be 'left' or 'right' for axis=1 and 'up' or 'down' for axis=0.    """    if invalid_val is np.nan:        mask = ~np.isnan(a)    else:        mask = a!=invalid_val    justified_mask = np.sort(mask,axis=axis)    if (side=='up') | (side=='left'):        justified_mask = np.flip(justified_mask,axis=axis)    out = np.full(a.shape, invalid_val)     if axis==1:        out[justified_mask] = a[mask]    else:        out.T[justified_mask.T] = a.T[mask.T]    return out

Sample runs -

In [473]: a # input arrayOut[473]: array([[1, 0, 2, 0],       [3, 0, 4, 0],       [5, 0, 6, 0],       [6, 7, 0, 8]])In [474]: justify(a, axis=0, side='up')Out[474]: array([[1, 7, 2, 8],       [3, 0, 4, 0],       [5, 0, 6, 0],       [6, 0, 0, 0]])In [475]: justify(a, axis=0, side='down')Out[475]: array([[1, 0, 0, 0],       [3, 0, 2, 0],       [5, 0, 4, 0],       [6, 7, 6, 8]])In [476]: justify(a, axis=1, side='left')Out[476]: array([[1, 2, 0, 0],       [3, 4, 0, 0],       [5, 6, 0, 0],       [6, 7, 8, 0]])In [477]: justify(a, axis=1, side='right')Out[477]: array([[0, 0, 1, 2],       [0, 0, 3, 4],       [0, 0, 5, 6],       [0, 6, 7, 8]])

Generic case (ndarray)

For a ndarray, we could modify it to -

def justify_nd(a, invalid_val, axis, side):        """    Justify ndarray for the valid elements (that are not invalid_val).    Parameters    ----------    A : ndarray        Input array to be justified    invalid_val : scalar        invalid value    axis : int        Axis along which justification is to be made    side : str        Direction of justification. Must be 'front' or 'end'.        So, with 'front', valid elements are pushed to the front and        with 'end' valid elements are pushed to the end along specified axis.    """        pushax = lambda a: np.moveaxis(a, axis, -1)    if invalid_val is np.nan:        mask = ~np.isnan(a)    else:        mask = a!=invalid_val    justified_mask = np.sort(mask,axis=axis)        if side=='front':        justified_mask = np.flip(justified_mask,axis=axis)                out = np.full(a.shape, invalid_val)    if (axis==-1) or (axis==a.ndim-1):        out[justified_mask] = a[mask]    else:        pushax(out)[pushax(justified_mask)] = pushax(a)[pushax(mask)]    return out

Sample runs -

Input array :

In [87]: aOut[87]: array([[[54, 57,  0, 77],        [77,  0,  0, 31],        [46,  0,  0, 98],        [98, 22, 68, 75]],       [[49,  0,  0, 98],        [ 0, 47,  0, 87],        [82, 19,  0, 90],        [79, 89, 57, 74]],       [[ 0,  0,  0,  0],        [29,  0,  0, 49],        [42, 75,  0, 67],        [42, 41, 84, 33]],       [[ 0,  0,  0, 38],        [44, 10,  0,  0],        [63,  0,  0,  0],        [89, 14,  0,  0]]])

To 'front', along axis =0 :

In [88]: justify_nd(a, invalid_val=0, axis=0, side='front')Out[88]: array([[[54, 57,  0, 77],        [77, 47,  0, 31],        [46, 19,  0, 98],        [98, 22, 68, 75]],       [[49,  0,  0, 98],        [29, 10,  0, 87],        [82, 75,  0, 90],        [79, 89, 57, 74]],       [[ 0,  0,  0, 38],        [44,  0,  0, 49],        [42,  0,  0, 67],        [42, 41, 84, 33]],       [[ 0,  0,  0,  0],        [ 0,  0,  0,  0],        [63,  0,  0,  0],        [89, 14,  0,  0]]])

Along axis=1 :

In [89]: justify_nd(a, invalid_val=0, axis=1, side='front')Out[89]: array([[[54, 57, 68, 77],        [77, 22,  0, 31],        [46,  0,  0, 98],        [98,  0,  0, 75]],       [[49, 47, 57, 98],        [82, 19,  0, 87],        [79, 89,  0, 90],        [ 0,  0,  0, 74]],       [[29, 75, 84, 49],        [42, 41,  0, 67],        [42,  0,  0, 33],        [ 0,  0,  0,  0]],       [[44, 10,  0, 38],        [63, 14,  0,  0],        [89,  0,  0,  0],        [ 0,  0,  0,  0]]])

Along axis=2 :

In [90]: justify_nd(a, invalid_val=0, axis=2, side='front')Out[90]: array([[[54, 57, 77,  0],        [77, 31,  0,  0],        [46, 98,  0,  0],        [98, 22, 68, 75]],       [[49, 98,  0,  0],        [47, 87,  0,  0],        [82, 19, 90,  0],        [79, 89, 57, 74]],       [[ 0,  0,  0,  0],        [29, 49,  0,  0],        [42, 75, 67,  0],        [42, 41, 84, 33]],       [[38,  0,  0,  0],        [44, 10,  0,  0],        [63,  0,  0,  0],        [89, 14,  0,  0]]])

To the 'end' :

In [94]: justify_nd(a, invalid_val=0, axis=2, side='end')Out[94]: array([[[ 0, 54, 57, 77],        [ 0,  0, 77, 31],        [ 0,  0, 46, 98],        [98, 22, 68, 75]],       [[ 0,  0, 49, 98],        [ 0,  0, 47, 87],        [ 0, 82, 19, 90],        [79, 89, 57, 74]],       [[ 0,  0,  0,  0],        [ 0,  0, 29, 49],        [ 0, 42, 75, 67],        [42, 41, 84, 33]],       [[ 0,  0,  0, 38],        [ 0,  0, 44, 10],        [ 0,  0,  0, 63],        [ 0,  0, 89, 14]]])

Thanks to all this is what I later use

def justify(a, direction):    mask = a>0    justified_mask = numpy.sort(mask,0) if direction == 'up' or direction =='down' else numpy.sort(mask, 1)    if direction == 'up':        justified_mask = justified_mask[::-1]    if direction =='left':        justified_mask = justified_mask[:,::-1]    if direction =='right':        justified_mask = justified_mask[::-1, :]        out = numpy.zeros_like(a)     out.T[justified_mask.T] = a.T[mask.T]    return out