如何计算含NaN的3D NumPy数组各子集列均值并替换NaN？

你的代码报错原因很明确：计算得到的列均值形状是(2,3)，而原数组是(2,3,3)，两者维度不匹配，无法完成广播运算。

修正方案1：用np.nanmean直接计算（更简洁） #

np.nanmean可以直接忽略NaN值计算均值，加上keepdims=True保持维度，让结果能和原数组匹配广播：

import numpy as np

a = np.array([[[1,2,3],[4,np.nan,6], [7,8,9]],
             [[11,12,13],[14,np.nan,16], [17,18,19]]])

# 计算每组的列均值，keepdims保持维度为(2,1,3)
col_means = np.nanmean(a, axis=1, keepdims=True)
# 替换NaN值
a_filled = np.where(np.isnan(a), col_means, a)

print(a_filled)

输出结果：

[[[ 1.  2.  3.]
  [ 4.  5.  6.]
  [ 7.  8.  9.]]

 [[11. 12. 13.]
  [14. 15. 16.]
  [17. 18. 19.]]]

修正方案2：保留掩码数组的写法 #

如果一定要用numpy.ma，只需要在计算均值时加上keepdims=True，让结果维度匹配：

import numpy as np
import numpy.ma as ma

a = np.array([[[1,2,3],[4,np.nan,6], [7,8,9]],
             [[11,12,13],[14,np.nan,16], [17,18,19]]])

# 计算列均值并保持维度为(2,1,3)
col_means = ma.array(a, mask=np.isnan(a)).mean(axis=1, keepdims=True)
# 替换NaN
a_filled = np.where(np.isnan(a), col_means, a)

print(a_filled)

核心逻辑：keepdims=True会把mean(axis=1)的结果从(2,3)变成(2,1,3)，这样第二个维度的1可以被广播到3，和原数组的每一行对应，解决维度不匹配的问题。

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