基于Numpy数组的条件计算优化：np.where为何未提速？

我拥有三个numpy数组，希望基于第一个数组的数值对另外两个数组进行计算。具体数组如下：

import numpy as np

a = np.array([1, 2, 3, 4, 5])
b = np.array([5, 4, 3, 2, 1])
c = np.array([0, 1, 2, 3, 4])

期望得到一个新数组：当a中的元素小于等于3时，取b与c的差值；当a中的元素大于3时，取b与c的和。

使用循环可以轻松实现，但我需要更快的代码。已尝试使用np.where，但速度并未提升：

import timeit

t = timeit.default_timer()
for _ in range(10000):
    d1 = np.zeros(5)
    for i in range(5):
        if a[i] <= 3:
            d1[i] = b[i] - c[i]
        else:
            d1[i] = b[i] + c[i]
print(f"Time: {timeit.default_timer() - t} s")  # prints: 0.025416199998289812 s

t = timeit.default_timer()
for _ in range(10000):
    d2 = np.where(a <= 3, b - c, b + c)
print(f"Time: {timeit.default_timer() - t} s")  # prints: 0.02637680000043474 s

请问我是否错误使用了np.where，或是有其他方法能让这段代码更快？

优化方法 #

你没有错误使用np.where，但在小数组场景下，Numpy的向量化操作优势不明显——因为Numpy函数调用存在固定开销，反而可能比简单循环慢。当数组规模扩大后，向量化方法的速度会远超循环。

如果要进一步优化当前场景的性能，可以试试预计算掩码+索引赋值的方式，避免np.where中同时计算b-c和b+c两个完整数组（即使只有部分结果会被用到）：

mask = a <= 3
d = b.copy()
d[mask] -= c[mask]
d[~mask] += c[~mask]

测试该方法的速度：

t = timeit.default_timer()
for _ in range(10000):
    mask = a <= 3
    d3 = b.copy()
    d3[mask] -= c[mask]
    d3[~mask] += c[~mask]
print(f"Time: {timeit.default_timer() - t} s")

这个方法的优势在于：

• 仅复制一次数组，避免了b-c和b+c两个完整数组的冗余计算
• 直接对数组切片进行操作，内存访问更高效

另外需要注意：当数组规模扩大到比如10000个元素时，循环的速度会急剧下降，而np.where和掩码赋值的向量化方法会保持稳定的高效性。

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