在Python中高效地进行二维地图扩散

Question

我刚开始学Python，所以在做一个项目。这个项目的一部分是要在地图上进行扩散。我实现这个功能的方法是，遍历每个格子，把当前格子的值设为它周围邻居（北、西、南、东）的值之和乘以0.2。如果我用C语言来做，我会用两个循环遍历一个数组，像这样：arr[i*宽度 + j] = arr[j+1]、arr[j-1]、arr[i+1]、arr[i-1]（就是邻居的值），然后对几个不同的数组做同样的操作（因为我会改变地图的不同特性）。但是，我不确定在Python中这样做是否是最快的。有些人建议我使用像numPy这样的库，但我的地图宽度大概不会超过200（最多4-5万个元素），我不确定使用这些库是否值得。我其实也不知道有什么内置函数可以实现我想要的功能。有没有什么建议呢？

补充：这个计算会非常复杂，也就是说每个位置都会有不简单的计算。

Answer 1

这个对unutbu代码的修改保持了数组的全局总和不变，同时又让数组中的值发生了变化：

        import numpy as np

        def diffuse(arr, d):
            contrib = (arr * d)
            w = contrib / 8.0
            r = arr - contrib
            N = np.roll(w, shift=-1, axis=0)
            S = np.roll(w, shift=1, axis=0)
            E = np.roll(w, shift=1, axis=1)
            W = np.roll(w, shift=-1, axis=1)
            NW = np.roll(N, shift=-1, axis=1)
            NE = np.roll(N, shift=1, axis=1)
            SW = np.roll(S, shift=-1, axis=1)
            SE = np.roll(S, shift=1, axis=1)
            diffused = r + N + S + E + W + NW + NE + SW + SE
            return diffused

Answer 2

使用卷积。

from numpy import *
from scipy.signal import convolve2d

mapArr=array(map)
kernel=array([[0  , 0.2,   0],
              [0.2,   0, 0.2],
              [0  , 0.2,   0]])
diffused=convolve2d(mapArr,kernel,boundary='wrap')

这是在说蚂蚁挑战吗？如果是的话，在蚂蚁的这个背景下，我的实现中，convolve2d的速度比循环快了大约20倍。

Answer 3

这件事用NumPy来处理非常简单。函数np.roll可以返回一个数组的副本，并且可以按照指定的方向“滚动”这个数组。

比如，给定一个数组x，

x=np.arange(9).reshape(3,3)
# array([[0, 1, 2],
#        [3, 4, 5],
#        [6, 7, 8]])

你可以用下面的方式把列向右滚动：

np.roll(x,shift=1,axis=1)
# array([[2, 0, 1],
#        [5, 3, 4],
#        [8, 6, 7]])

使用np.roll时，边界会像在一个环面上那样连接起来。如果你不想要这种连接的边界，可以在数组的边缘加上一圈零，然后在每次迭代之前把边缘的值重置为零。

import numpy as np

def diffusion(arr):
    while True:
        arr+=0.2*np.roll(arr,shift=1,axis=1) # right
        arr+=0.2*np.roll(arr,shift=-1,axis=1) # left
        arr+=0.2*np.roll(arr,shift=1,axis=0) # down
        arr+=0.2*np.roll(arr,shift=-1,axis=0) # up
        yield arr

N=5
initial=np.random.random((N,N))
for state in diffusion(initial):
    print(state)
    raw_input()