如何实现NumPy向量化分段三角函数并用matplotlib绘制图像？

核心问题说明 #

你目前遇到的实现卡点本质是：单条件的np.where无法覆盖多层分段规则，现有代码还存在缺失依赖导入、绘图接口调用错误、逻辑不匹配分段要求的问题，完全没有用到NumPy对数组逐元素判断的原生能力。

针对给定的分段规则 t(x) = 0, x < 0; x, 0 <= x <1; 2-x, 1 <= x < 2, 0, x >= 2，完全不需要写逐元素循环，靠NumPy的原生向量化判断就能完成适配。

两种可落地的向量化实现方式 #

• 嵌套np.where实现：从最外层的边界条件开始逐层嵌套判断，逻辑链路和分段规则一一对应
• 布尔掩码赋值实现：先初始化全零结果数组，再按分段区间的布尔条件逐段赋值，代码可读性更强

完整可运行代码（嵌套np.where版本） #

import numpy as np
import matplotlib.pyplot as plt

# 支持NumPy数组输入的向量化三角函数
def triangle(x):
    return np.where(
        x < 0,
        0,
        np.where(
            x < 1,
            x,
            np.where(
                x < 2,
                2 - x,
                0
            )
        )
    )

# 生成高密度采样点绘制平滑曲线
x = np.linspace(-1, 8, 1000)
t_x = triangle(x)

plt.plot(x, t_x, linewidth=2, label='t(x)')
# 坐标轴范围留出少量边距，避免曲线贴边
plt.xlim(-1, 8)
plt.ylim(-0.1, 1.1)
plt.grid(alpha=0.3)
plt.legend()
plt.show()

可读性更强的布尔掩码版本函数替换 #

如果觉得嵌套np.where的层级结构不好维护，可以用掩码赋值的写法替换triangle函数，效果完全一致：

def triangle(x):
    # 生成和输入形状一致的全零结果数组
    res = np.zeros_like(x, dtype=np.float64)
    # 匹配0<=x<1的区间，赋值为x本身
    mask_01 = (x >= 0) & (x < 1)
    res[mask_01] = x[mask_01]
    # 匹配1<=x<2的区间，赋值为2-x
    mask_12 = (x >= 1) & (x < 2)
    res[mask_12] = 2 - x[mask_12]
    return res

常见错误避坑 #

• 代码必须先导入numpy，否则np.where、数组操作都会触发未定义错误
• plt.show()方法不接收任何参数，绘图流程必须先生成x采样点、计算对应函数值、通过plt.plot()传入数据后，再调用plt.show()展示画布
• 不要用Python层的for循环遍历数组元素做分段判断，这种写法完全丢失了NumPy向量化操作的性能优势
• 绘制曲线时要生成足够密集的x采样点，否则曲线会出现明显的折线锯齿

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