Skip to main content
 首页 » 编程设计

python之如何使用 numpy 使 for 循环更快

2025年05月04日17del

如果输入一个 numpy 二维数组,例如

                           100   100   100   100   100 
                           100    0     0     0    100 
                           100    0     0     0    100 
                           100    0     0     0    100 
                           100   100   100   100   100 

应该得到这样的输出

                            100   100   100   100   100 
                            100    50    25    50   100 
                            100    25     0    25   100 
                            100    50    25    50   100 
                            100   100   100   100   100 

其中除边框外的每个数字都成为其相邻数字的平均值。

我当前的代码有效,但我需要在没有 for 循环的情况下使用它,并使用 numpy 对其进行矢量化。

我当前的代码:

import numpy as np 
def evolve_heat_slow(u):       
    u2 = np.copy(u) 
    x=u2.shape[0] 
    y=u2.shape[1] 
    for i in range(1,x-1): 
        for s in range(1,y-1): 
            u2[i,s]=(u[i-1,s]+u[i+1,s]+u[i,s+1]+u[i,s-1])/4 
    return u2 

请您参考如下方法:

这几乎就是二维卷积的定义。 scipy 为您提供帮助。我复制 a 以保留边框; valid 模式下的卷积将生成一个较小的数组(没有边框),然后我将其粘贴到准备好的“框架”中。

import numpy as np 
from scipy.signal import convolve2d 
 
a = np.array([[100, 100, 100, 100, 100], [100, 0, 0, 0, 100],  [100, 0, 0, 0, 100], [100, 0, 0, 0, 100], [100, 100, 100, 100, 100]]) 
b = np.array([[0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0]]) 
r = np.copy(a) 
r[1:-1, 1:-1] = convolve2d(a, b, mode='valid') 
r 
# => array([[100, 100, 100, 100, 100], 
#           [100,  50,  25,  50, 100], 
#           [100,  25,   0,  25, 100], 
#           [100,  50,  25,  50, 100], 
#           [100, 100, 100, 100, 100]])