NumPy 配列内の要素を効率的に位置揃えするにはどうすればよいですか?

def justify(a, invalid_val=0, axis=1, side='left'):
    justified_mask = np.sort(a!=invalid_val, axis=axis)
    if (side=='up') or (side=='left'):
        justified_mask = np.flip(justified_mask,axis=axis)
    out = np.full(a.shape, invalid_val)
    if axis==1:
        out[justified_mask] = a[a!=invalid_val]
    else:
        out.T[justified_mask.T] = a.T[a.T!=invalid_val]
    return out

a = np.array([[1, 0, 2, 0],
               [3, 0, 4, 0],
               [5, 0, 6, 0],
               [0, 7, 0, 8]])

print(justify(a, axis=0, side='up'))  # Justify values vertically "up"
print(justify(a, axis=0, side='down'))  # Justify values vertically "down"
print(justify(a, axis=1, side='left'))  # Justify values horizontally "left"
print(justify(a, axis=1, side='right'))  # Justify values horizontally "right"

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

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

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

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

def justify_nd(a, invalid_val, axis, side):
    justified_mask = np.sort(a!=invalid_val, axis=axis)
    if side=='front':
        justified_mask = np.flip(justified_mask,axis=axis)
    out = np.full(a.shape, invalid_val)
    pushax = lambda a: np.moveaxis(a, axis, -1)
    if (axis==-1) or (axis==a.ndim-1):
        out[justified_mask] = a[a!=invalid_val]
    else:
        pushax(out)[pushax(justified_mask)] = pushax(a)[pushax(a!=invalid_val)]
    return out

a = np.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]]])

print(justify_nd(a, invalid_val=0, axis=0, side='front'))  # Justify first dimension "front"
print(justify_nd(a, invalid_val=0, axis=1, side='front'))  # Justify second dimension "front"
print(justify_nd(a, invalid_val=0, axis=2, side='front'))  # Justify third dimension "front"
print(justify_nd(a, invalid_val=0, axis=2, side='end'))  # Justify third dimension "end"

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

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

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