为何Python中切片数组变体比逐元素操作运行更慢？

问题背景 #

我参考了一个已解决的问题，其中存在冗余操作，于是尝试用掩码索引的方式消除冗余，结果和原代码一致，但性能却不如预期：采用双重循环逐元素判断的Variant1，反而比结合切片与掩码的Variant2快约30%。原以为Numpy切片底层由C实现，性能必然优于Python手动循环，因此产生疑问：为什么会出现这种情况？切片操作在哪些场景下比逐元素操作更快？它们是不是徒有其表的语法糖？

运行输出 #

D:\python\animation>python test.py
used time for variant 1: 1.0377624034881592
used time for variant 2: 1.30381441116333

D:\python\animation>python test.py
used time for variant 1: 0.8954949378967285
used time for variant 2: 1.251044750213623

D:\python\animation>python test.py
used time for variant 1: 0.9750621318817139
used time for variant 2: 1.3896379470825195

测试代码 #

import numpy as np
import numpy.ma as ma
import time


def test():

    f = np.array([
        [0,   0,   0,   0,   0,   0,      0], 
        [0,   1,   3,   6 ,  4,   2,      0], 
        [0,   2,   4,   7 ,  6,   4,      0],    
        [0,   0,   0,   0,   0,   0,      0]
        ])
        

    u = np.array([
        [0,   0,    0,    0,   0,   0,     0], 
        [0,   0.5,  1,    0,  -1,  -0.5,   0], 
        [0,   0.7,  1.1,  0,  -1,  -0.4,   0], 
        [0,   0,    0,    0,   0,   0,     0], 
        ])
        
    
    # calculate : variant 1
    x = np.zeros_like(f)
    
    maxcount = 100000
    
    start = time.time()

    for count in range(maxcount):
        for i in range(1,u.shape[0]-1):
            for j in range(1,u.shape[1]-1):
                if u[i,j] > 0: 
                    x[i,j] = u[i,j]*(f[i,j]-f[i,j-1])
                else:
                    x[i,j] = u[i,j]*(f[i,j+1]-f[i,j])
                
    end = time.time()
    print("used time for variant 1:", end-start)
                
                            
    
    # calculate : variant 2

    y = np.zeros_like(f)  

    
    start = time.time()
    
    for count in range(maxcount):
        maskl = (u[1:-1, 1:-1] > 0)
        maskr = ~maskl 
        diff = f[1:-1, 1:]  - f[1:-1, 0:-1]
        
        (y[1:-1, 1:-1])[maskl]  = (u[1:-1, 1:-1 ])[maskl]  * (diff[:, :-1])[maskl]
        (y[1:-1, 1:-1])[maskr]  = (u[1:-1, 1:-1 ])[maskr]  * (diff[:, 1: ])[maskr]
    
    end = time.time()
    print("used time for variant 2:", end-start)
    
    np.testing.assert_array_equal(x, y)


test()

小幅优化后的代码片段 #

提前获取切片能略微提升性能，但效果有限：

for count in range(maxcount):
    maskl = (u[1:-1, 1:-1] > 0)
    maskr = ~maskl 
    diff = f[1:-1, 1:]  - f[1:-1, 0:-1]
    
    yy = (y[1:-1, 1:-1])   # <<-- 
    uu = (u[1:-1, 1:-1 ])  # <<--
    
    yy[maskl]  = uu[maskl]  * (diff[:, :-1])[maskl]
    yy[maskr]  = uu[maskr]  * (diff[:, 1: ])[maskr]

原因分析 #

1. 小数据量下的开销占比过高 #

你的测试数组核心计算区域仅为2x4，数据量极小。Numpy的切片、掩码操作需要在Python层和C层之间交互，还要创建临时数组（比如maskl、diff），这些额外开销在小数据量下占比极高，反而盖过了C层计算的优势。而Python循环虽然是解释执行，但数据量小意味着循环次数少，开销可以忽略。

2. 掩码索引的内存访问劣势 #

掩码索引选取的元素是非连续的，会导致内存缓存命中率降低；而Python双重循环是连续访问数组元素，缓存友好性更好，执行效率更高。

3. 循环内的重复冗余计算 #

Variant2在每次maxcount循环中都要重新计算掩码、差值数组，这些都是不必要的重复操作；而Variant1的循环内仅做简单的条件判断和算术计算，没有额外的数组创建开销。

切片操作的优势场景 #

切片并非只是语法糖，在以下场景中性能会远超Python逐元素循环：

• 大数据量计算：当数组规模达到千级以上时，Numpy向量化操作能把计算完全放在C层执行，避免Python解释器的循环开销，性能优势会呈指数级放大。
• 连续内存批量操作：比如对整个数组做算术运算、连续切片的赋值/计算，这类操作是连续内存访问，缓存命中率高，能充分发挥Numpy的底层优化。
• 高频循环场景：当外层循环次数极多（比如百万级以上），Python解释器的循环开销会被无限放大，此时向量化操作的优势会非常明显。

优化建议 #

要提升Variant2的性能，关键是减少循环内的重复操作，改用更高效的向量化逻辑：

# 优化后的Variant2：将固定计算移到循环外，用np.where简化逻辑
y = np.zeros_like(f)
# 提前计算所有需要的切片和差值，避免循环内重复计算
core_u = u[1:-1, 1:-1]
core_y = y[1:-1, 1:-1]
diff_left = f[1:-1, 1:-1] - f[1:-1, :-2]
diff_right = f[1:-1, 2:] - f[1:-1, 1:-1]

start = time.time()
for count in range(maxcount):
    core_y[:] = np.where(core_u > 0, core_u * diff_left, core_u * diff_right)
end = time.time()
print("used time for optimized variant 2:", end-start)

这种优化方式减少了循环内的临时数组创建，性能会接近甚至超过Variant1，同时保持了向量化代码的简洁性。

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