这是一种受this other post 启发的矢量化方法,并泛化为涵盖所有四个方向的non-zeros -
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)
对于一个 ndarray,我们可以将其修改为 -
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
示例运行 -
输入数组:
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]]])
致'front',沿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]]])
沿着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]]])
沿着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]]])
致'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]]])