使用NumPy掩码数组处理嵌套数组时避免输出扁平化的问题

错误原因 #

你的代码核心问题在于布尔数组直接索引numpy数组会触发扁平化操作：

• 当你用x[bool_arr]时，numpy会把所有满足条件的元素从二维数组中提取出来，直接拼成一个一维数组，完全丢失了原有的行结构。
• 另外，原数组x因为混合了整数和字符串'NaN'，numpy会自动将整个数组转换为字符串类型，但这不是导致结构丢失的直接原因。

修正方案 #

要保留每行的嵌套结构，需要逐行处理过滤，而不是用全局的布尔索引。以下是两种可行的实现方式：

方法1：列表推导式逐行过滤 #

import numpy as np

x = np.array([[1, 2, 3], ['NaN', 4, 'NaN'], [7, 8, 9]])
# 遍历每一行，过滤掉'NaN'元素并转为列表
filtered_arr = [row[row != 'NaN'].tolist() for row in x]
print(filtered_arr)

输出结果：

[['1', '2', '3'], ['4'], ['7', '8', '9']]

方法2：使用numpy的apply_along_axis逐行处理 #

import numpy as np

x = np.array([[1, 2, 3], ['NaN', 4, 'NaN'], [7, 8, 9]])

# 定义单行过滤函数
def filter_nan(row):
    return row[row != 'NaN'].tolist()

# 沿轴1（行方向）应用过滤函数
filtered_arr = np.apply_along_axis(filter_nan, axis=1, arr=x).tolist()
print(filtered_arr)

输出结果和方法1一致。

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