如何为numpy数组编写依输入数值分支的高效无显式循环函数？

针对你需要实现的分段函数：

f(x) = 1, if x<1
= g(x), for 1<=x<10
= 0, for x >=10

当输入是Numpy数组时，完全不需要显式循环，用Numpy的向量化特性就能高效实现，下面是两种常用方案：

方案一：布尔索引（直观易读） #

利用Numpy的布尔掩码直接给数组对应位置赋值，底层是C级别的向量化操作，效率极高：

import numpy as np

# 示例g(x)，可替换为任意支持Numpy数组的简单函数
def g(x):
    return x ** 2

def f(x):
    # 先初始化结果数组为0，对应x>=10的情况
    result = np.zeros_like(x)
    # 给x<1的元素赋值1
    result[x < 1] = 1
    # 给1<=x<10的元素赋值g(x)的计算结果
    middle_mask = (1 <= x) & (x < 10)
    result[middle_mask] = g(x[middle_mask])
    return result

方案二：np.piecewise（专为分段函数设计） #

Numpy内置的np.piecewise函数可以直接定义分段逻辑，代码更紧凑：

import numpy as np

def g(x):
    return x ** 2

def f(x):
    return np.piecewise(
        x,
        # 按顺序定义三个分段条件
        [x < 1, (1 <= x) & (x < 10), x >= 10],
        # 对应每个条件的输出：常数或函数
        [1, g, 0]
    )

注意事项 #

• 两种方案都避免了显式循环，性能远高于遍历数组元素；
• 如果你的g(x)不是原生支持Numpy数组的函数，需要先将其改造为向量化版本（比如用Numpy内置函数实现），尽量不要用np.vectorize——它本质是包装了循环，效率提升有限。

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