基于Numpy的嵌套for循环向量化及等值线绘制优化

一、向量化改造方案（完全替代嵌套循环） #

你的嵌套循环效率低是因为逐元素遍历，numpy的核心优势就是向量化运算，直接对整个网格数组操作，不需要循环。下面是改造后的代码：

import numpy as np

rmin = -5.0
rmax = 5.0
c = 4.0
x = np.arange(rmin, rmax, 0.1)
y = np.arange(rmin, rmax, 0.1)
x, y = np.meshgrid(x, y)
f = lambda x,y: y**2.0 - 4*x

# 向量化计算函数值与筛选条件
z = f(x, y)
mask = np.abs(z - c) < 1e-4
# 提取符合条件的坐标并整理成目标格式
realplots = np.column_stack((x[mask], y[mask])).tolist()

关键说明： #

• 直接对二维的x、y数组调用f，numpy会自动进行广播运算，一次性计算所有点的函数值，比循环快几个数量级。
• 用布尔掩码mask筛选出满足条件的点，再通过np.column_stack把x、y的筛选结果组合成坐标对，最后转成列表格式和原代码输出一致。

解决你之前的np.where报错问题： #

你之前的错误是误用了np.where的返回形式，正确的用法是先获取满足条件的索引，再提取坐标：

rows, cols = np.where(mask)
realplots = [[x[r, c], y[r, c]] for r, c in zip(rows, cols)]

np.where返回的是行、列索引的数组，不能直接返回[0,0]这种标量对，因为形状和输入不匹配，上面的写法就能避免这个问题。

二、更高效的等值线绘制方法（无需手动筛选点） #

如果你的最终目的是绘制等值线，完全不需要手动收集离散点，用matplotlib的contour函数是最优解——它专门处理等值线绘制，内部实现更高效，还能生成连续平滑的线条：

import numpy as np
import matplotlib.pyplot as plt

rmin = -5.0
rmax = 5.0
c = 4.0
x = np.arange(rmin, rmax, 0.1)
y = np.arange(rmin, rmax, 0.1)
x, y = np.meshgrid(x, y)
f = lambda x,y: y**2.0 - 4*x

z = f(x, y)
# 绘制c值对应的等值线
plt.contour(x, y, z, levels=[c], colors='red')
plt.xlabel('x')
plt.ylabel('y')
plt.show()

优势： #

• 不管f怎么动态修改，只要重新计算z = f(x, y)，再调用contour就能更新等值线，完全不依赖f的具体数学形式。
• 生成的是连续线条，比离散点连接的线条更美观，计算效率也远高于手动筛选点。

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