能否在Matplotlib的scatter()散点图中添加多项式回归线？

当然可以！给散点图添加多项式回归线其实很简单，我们可以借助NumPy的多项式拟合工具，再结合Matplotlib画出平滑的拟合曲线。下面结合你提供的代码和数据来一步步实现：

步骤1：导入必要的库 #

除了你已经用到的matplotlib.pyplot，还需要导入numpy来完成多项式拟合：

import numpy as np

步骤2：执行多项式拟合 #

我们可以用np.polyfit()函数计算多项式的系数，再用np.poly1d()生成对应的多项式函数。这里以3次多项式为例（你可以根据数据的分布调整degree参数，次数越高拟合越贴近原始数据，但要注意避免过拟合）：

# 选择多项式次数
degree = 3

# 分别拟合两组数据的回归线
# Orig HC的拟合
orig_coeffs = np.polyfit(alg_n, orig_hc_runtime, degree)
orig_poly = np.poly1d(orig_coeffs)

# Mod HC的拟合
mod_coeffs = np.polyfit(alg_n, mod_hc_runtime, degree)
mod_poly = np.poly1d(mod_coeffs)

步骤3：生成平滑的拟合曲线x值 #

为了让回归线看起来更平滑，我们需要生成一组连续的x值（而不是用原始离散的alg_n）：

x_fit = np.linspace(min(alg_n), max(alg_n), 100)

步骤4：绘制多项式回归线 #

用plt.plot()把拟合曲线画出来，颜色和标签和对应的散点保持一致：

plt.plot(x_fit, orig_poly(x_fit), color="b", linestyle="-", label="Orig HC Polynomial Fit")
plt.plot(x_fit, mod_poly(x_fit), color="c", linestyle="-", label="Mod HC Polynomial Fit")

完整示例代码 #

把以上步骤整合到你的原有代码中，最终代码如下：

import matplotlib.pyplot as plt
import numpy as np

# 你的原始数据
alg_n = [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]
orig_hc_runtime = [0.01, 0.02, 0.03, 0.04, 0.04, 0.04, 0.05, 0.09, 0.1, 0.12, 0.13, 0.15]
mod_hc_runtime = [0.008, 0.015, 0.02, 0.03, 0.03, 0.035, 0.04, 0.07, 0.08, 0.1, 0.11, 0.12]

# 绘制散点图
plt.scatter(alg_n, orig_hc_runtime, label="Orig HC", color="b", s=4)
plt.scatter(alg_n, mod_hc_runtime, label="Mod HC", color="c", s=4)

# 多项式拟合
degree = 3
orig_coeffs = np.polyfit(alg_n, orig_hc_runtime, degree)
orig_poly = np.poly1d(orig_coeffs)
mod_coeffs = np.polyfit(alg_n, mod_hc_runtime, degree)
mod_poly = np.poly1d(mod_coeffs)

# 生成平滑x值
x_fit = np.linspace(min(alg_n), max(alg_n), 100)

# 绘制回归线
plt.plot(x_fit, orig_poly(x_fit), color="b", linestyle="-", label="Orig HC Polynomial Fit")
plt.plot(x_fit, mod_poly(x_fit), color="c", linestyle="-", label="Mod HC Polynomial Fit")

# 你的原有绘图设置
x_values = [x for x in range(5, len(alg_n)+2, 2)]  # 根据实际n_init调整
y_values = [y for y in range(0, 10, 2)]
plt.xlabel("Number of Queues")
plt.ylabel("Runtime")
plt.legend()
plt.show()

小提示 #

• 如果你不确定选多少次多项式，可以尝试不同的degree值，观察拟合效果；
• 如果数据存在明显的趋势，低次多项式（如2次）可能就足够，高次容易出现过拟合；
• 你也可以用np.polyval(orig_coeffs, x_fit)来计算拟合的y值，效果和orig_poly(x_fit)一致。

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