编写忽略NaN值，对m×n数组执行numpy.polyval的函数

你的思路是对的：遍历数组处理非NaN值，保留原NaN位置，但当前代码存在几个可以优化的点，我来帮你梳理并给出改进方案：

原代码的问题分析 #

1. 浅拷贝导致原数组被修改：data2 = data只是创建了原数组的引用，不是独立副本。这意味着你在修改data2时，输入的data数组也会被改变——这几乎不是你想要的行为。
2. 循环索引错误：np.arange(1,m)和np.arange(1,n)会跳过第0行和第0列的元素，这些位置的非NaN值完全没被处理。
3. 效率低下：嵌套的for循环在处理大型数组时速度很慢，numpy的核心优势是向量化操作，应该尽量避免显式循环。

方案1：修复循环的基础版本 #

如果你想保留循环的逻辑，先修正上述问题：

import numpy as np

def polyval_NaN(p, data):
    # 创建输入数组的独立副本，彻底隔离原数据
    data2 = data.copy()
    m, n = data.shape  # 把shape获取放在函数开头，避免重复计算
    # 从0开始遍历所有行和列
    for i in np.arange(m):
        for j in np.arange(n):
            if not np.isnan(data2[i, j]):
                data2[i, j] = np.polyval(p, data[i, j])
    return data2

方案2：高效向量化版本（推荐） #

利用numpy的布尔索引实现批量操作，代码更简洁，速度快得多（尤其大数组）：

import numpy as np

def polyval_NaN(p, data):
    # 先创建副本保护原数据
    data2 = data.copy()
    # 生成非NaN元素的掩码
    non_nan_mask = ~np.isnan(data2)
    # 对所有非NaN元素批量应用polyval
    data2[non_nan_mask] = np.polyval(p, data2[non_nan_mask])
    return data2

甚至可以简化成更紧凑的写法：

import numpy as np

def polyval_NaN(p, data):
    data2 = data.copy()
    data2[~np.isnan(data2)] = np.polyval(p, data2[~np.isnan(data2)])
    return data2

测试示例 #

用一个简单的测试数据验证效果：

# 定义多项式：2x² + 3x + 1
p = [2, 3, 1]
test_data = np.array([[1, np.nan, 3], [np.nan, 5, 2]])

result = polyval_NaN(p, test_data)
print(result)

输出结果：

[[ 6. nan 22.]
 [nan 56. 15.]]

可以看到非NaN值都正确计算了，NaN位置也保留了。

总结建议 #

优先使用向量化版本，numpy的底层是C实现的向量化操作，比Python级别的for循环效率高几个数量级。同时一定要记得创建数组副本，避免意外修改原始输入数据。

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