pythonnumpysearchartificial-intelligence2048

Python: Justifying NumPy array


Please I am a bit new to Python and it has been nice, I could comment that python is very sexy till I needed to shift content of a 4x4 matrix which I want to use in building a 2048 game demo of the game is here I have this function

def cover_left(matrix):
        new=[[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]]
        for i in range(4):
             count=0
             for j in range(4):
                if mat[i][j]!=0:
                    new[i][count]=mat[i][j]
                    count+=1
        return new

This is what this function does if you call it like this

cover_left([
              [1,0,2,0], 
              [3,0,4,0], 
              [5,0,6,0], 
              [0,7,0,8]
          ])

It will cover the zeros to the left and produce

[  [1, 2, 0, 0],
   [3, 4, 0, 0],
   [5, 6, 0, 0],
   [7, 8, 0, 0]]

Please I need someone to help me with a numpy way of doing this which I believe will be faster and require less code (I am using in a depth-first search algo) and more importantly the implementation of cover_up, cover_down and cover_left.

`cover_up`
    [  [1, 7, 2, 8],
       [3, 0, 4, 0],
       [5, 0, 6, 0],
       [0, 0, 0, 0]]
`cover_down`
    [  [0, 0, 0, 0],
       [1, 0, 2, 0],
       [3, 0, 4, 0],
       [5, 7, 6, 8]]
`cover_right`
    [  [0, 0, 1, 2],
       [0, 0, 3, 4],
       [0, 0, 5, 6],
       [0, 0, 7, 8]]

Solution

  • 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 array
    Out[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]: a
    Out[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]]])