NumPy 数组如何有效地对齐（移位）？

NumPy 数组的合理性

问题：

优化用于在 2048 游戏演示中使用矩阵中内容的代码。具体来说，需要函数将矩阵中的非零值向左、向右、向上或向下移动。

使用 NumPy 的解决方案：

提供的代码提供了一种矢量化方法，其灵感来自另一篇文章：

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

示例运行：

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

通用案例(ndarray):

对于通用的n维数组，代码可以修改如下：

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

示例运行 (ndarray):

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

到“前面”，沿轴=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]]])

沿轴=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]]])

沿轴=2 :

以上是NumPy 数组如何有效地对齐（移位）？的详细内容。更多信息请关注PHP中文网其他相关文章！