如何过滤3D数组中子2D数组第0列为NaN的行？

原代码的核心问题 #

1. NaN判断逻辑错误：不能用!= np.nan检测NaN，因为NaN的特性是「不等于任何值，包括自身」，所以arr[:, :, 0] != np.nan会返回全True的数组，完全起不到过滤作用，必须用np.isnan()函数来识别NaN。
2. 维度结构丢失：np.where返回的是一维索引，直接用arr[np.where(...)]会把符合条件的元素拉平，最终得到2D数组，破坏了原3D层级。

方案1：列表推导式（推荐，适配行数不一致的情况） #

遍历每个子2D数组，单独过滤第0列非NaN的行，结果会是一个保留原层级的2D数组列表，天然适配各子数组过滤后行数不同的场景：

import numpy as np

# 构造你的3D数组
arr = np.array([
    [[1., 0., 0., 10., 0.],
     [2., 0., 0., 9., 0.],
     [np.nan, 0., 0., 8., 0.],
     [4., 0., 0., 7., 0.],
     [np.nan, 0., 0., 6., 0.]],
    [[1., 0., 0., 199., 0.],
     [2., 0., 0., 198., 0.],
     [3., 0., 0., 196., 0.],
     [np.nan, 0., 0., 190., 0.],
     [np.nan, 0., 0., 160., 0.]],
    [[1., 0., 0., 999., 0.],
     [2., 0., 0., 870., 0.],
     [np.nan, 0., 0., 270., 0.],
     [np.nan, 0., 0., 100., 0.],
     [np.nan, 0., 0., 80., 0.]]
])

# 过滤每个子数组中第0列非NaN的行
filtered_list = [sub_arr[~np.isnan(sub_arr[:, 0])] for sub_arr in arr]

# 查看结果
for idx, sub in enumerate(filtered_list):
    print(f"第{idx+1}个子数组过滤后形状：{sub.shape}")
    print(sub)

输出示例：

第1个子数组过滤后形状：(3, 5)
[[ 1.  0.  0. 10.  0.]
 [ 2.  0.  0.  9.  0.]
 [ 4.  0.  0.  7.  0.]]
第2个子数组过滤后形状：(3, 5)
[[  1.   0.   0. 199.   0.]
 [  2.   0.   0. 198.   0.]
 [  3.   0.   0. 196.   0.]]
第3个子数组过滤后形状：(2, 5)
[[  1.   0.   0. 999.   0.]
 [  2.   0.   0. 870.   0.]]

方案2：布尔索引（仅适用于所有子数组过滤后行数相同的情况） #

如果能保证每个子数组过滤后保留的行数一致，可以用布尔索引结合形状重塑，得到标准的3D数组：

# 生成每行的有效mask：(3,5)的布尔数组
mask = ~np.isnan(arr[:, :, 0])

# 过滤后重塑为3D数组（假设每个子数组保留3行）
filtered_arr = arr[mask].reshape(arr.shape[0], -1, arr.shape[2])
print(filtered_arr.shape)  # 输出 (3, 3, 5)

内容的提问来源于stack exchange，提问作者Taz