我有一个多维数组。示例(2D):
x = np.array([[ 1., 1., np.nan, np.nan],
[ 2., np.nan, 2., np.nan],
[ np.nan, 3., np.nan, np.nan]])
是否有一种简单有效的方法可以沿着轴“压缩”/“挤压”/“推”出它的nans?我的意思是,输出(这里:axis = 0)将成为:
np.array([[ 1., 1., np.nan, np.nan],
[ 2., 3., 2., np.nan]])
还应该使用2个以上的维度。
你可以在非argsort
元素的面具上使用nan
;使用稳定的排序算法(如mergesort)来保留非nan
元素的原始顺序:
mask = np.isnan(x)
cut = np.min(np.count_nonzero(mask, axis=0))
x[np.argsort(~mask, axis=0, kind='mergesort')[cut:], np.arange(x.shape[1])]
输出:
array([[ 1., 1., nan, nan],
[ 2., 3., 2., nan]])
ND-版本:
import numpy as np
def nan_bouncer(x, axis=0):
if axis != 0:
x = np.moveaxis(x, axis, 0)
mask = np.isnan(x)
cut = np.min(np.count_nonzero(mask, axis=0))
idx = tuple(np.ogrid[tuple(map(slice, x.shape[1:]))])
res = x[(np.argsort(~mask, axis=0, kind='mergesort')[cut:],) + idx]
return res if axis == 0 else np.moveaxis(res, 0, axis)
#demo
data = np.random.randint(0, 3, (3, 4, 4)).astype(float)
data /= data / data
print(data)
print(nan_bouncer(data))
print(nan_bouncer(data, 2))
样本输出:
[[[ nan 1. 2. 1.]
[ 2. nan nan 2.]
[ 2. 1. 1. 2.]
[ 1. 1. 2. nan]]
[[ nan nan 2. 1.]
[ 2. 2. nan 1.]
[ 2. 2. 2. 2.]
[ 2. 2. nan 1.]]
[[ 1. 1. nan nan]
[ 1. 1. 2. 1.]
[ 2. nan 2. 1.]
[ 1. 1. 1. 2.]]]
[[[ nan nan nan nan]
[ 2. nan nan 2.]
[ 2. nan 1. 2.]
[ 1. 1. nan nan]]
[[ nan 1. 2. 1.]
[ 2. 2. nan 1.]
[ 2. 1. 2. 2.]
[ 2. 2. 2. 1.]]
[[ 1. 1. 2. 1.]
[ 1. 1. 2. 1.]
[ 2. 2. 2. 1.]
[ 1. 1. 1. 2.]]]
[[[ nan 1. 2. 1.]
[ nan nan 2. 2.]
[ 2. 1. 1. 2.]
[ nan 1. 1. 2.]]
[[ nan nan 2. 1.]
[ nan 2. 2. 1.]
[ 2. 2. 2. 2.]
[ nan 2. 2. 1.]]
[[ nan nan 1. 1.]
[ 1. 1. 2. 1.]
[ nan 2. 2. 1.]
[ 1. 1. 1. 2.]]]