含NaN值数组对应位置求平均如何纳入仅单个数组的有效值

需求明确 #

计算两个等长数值数组A、B的逐位置平均值，规则如下：

• 同位置两元素均为有效数值：返回两数算术平均值
• 同位置仅单个元素为有效数值：直接返回该有效数值
• 同位置两元素均为NaN：返回NaN

原有逻辑(A + B) / 2的问题是：只要任意一个位置元素为NaN，加法结果就为NaN，导致单有效值的位置无法得到正确结果。

最优实现（NumPy内置方法） #

直接使用NumPy自带的np.nanmean方法，该方法会自动忽略NaN值计算均值，完全匹配上述规则，代码简洁且运行效率高：

import numpy as np

# 将两个数组沿新轴堆叠为二维数组，沿轴0计算逐位置忽略NaN的均值
avg = np.nanmean(np.stack([A, B], axis=0), axis=0)

逻辑匹配验证 #

对应提供的样例数据：

• 索引10、11位置：仅A存在有效值，计算结果直接返回A对应位置的109.93013333、121.27613333
• 索引12~27位置：两数组均有有效值，计算结果和现有正确的avg数组完全一致
• 索引28位置：仅B存在有效值82.5104，计算结果直接返回该值
• 其余两数组均为NaN的位置：返回NaN

手动分支实现（自定义逻辑场景） #

如果需要手动写判断逻辑，可以通过掩码分情况赋值，结果和上述方法完全一致：

import numpy as np

# 生成两个数组的有效值掩码
a_valid = ~np.isnan(A)
b_valid = ~np.isnan(B)

# 分情况计算
avg = np.full_like(A, np.nan)
avg[a_valid & b_valid] = (A[a_valid & b_valid] + B[a_valid & b_valid]) / 2
avg[a_valid & ~b_valid] = A[a_valid & ~b_valid]
avg[~a_valid & b_valid] = B[~a_valid & b_valid]

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