Question: I want to avoid using for loop in the following code to achieve performance. Is vectorization suitable for this kind of problem?

```
a = np.array([[0,1,2,3,4],
              [5,6,7,8,9],
              [0,1,2,3,4],
              [5,6,7,8,9],
              [0,1,2,3,4]],dtype= np.float32)

temp_a = np.copy(a)

for i in range(1,a.shape[0]-1):
    for j in range(1,a.shape[1]-1):
        if a[i,j] > 3:
            temp_a[i+1,j] += a[i,j] / 5.
            temp_a[i-1,j] += a[i,j] / 5.
            temp_a[i,j+1] += a[i,j] / 5.
            temp_a[i,j-1] += a[i,j] / 5.
            temp_a[i,j]   -= a[i,j] * 4. / 5.
a = np.copy(temp_a)   

```

Answer: 
You are basically doing convolution, with some special treatment for borders.

Try the following:

```
from scipy.signal import convolve2d

# define your filter
f = np.array([[0.0, 0.2, 0.0],
              [0.2,-0.8, 0.2],
              [0.0, 0.2, 0.0]])

# select parts of 'a' to be used for convolution
b = (a * (a > 3))[1:-1, 1:-1]

# convolve, padding with zeros ('same' mode)
c = convolve2d(b, f, mode='same')

# add the convolved result to 'a', excluding borders
a[1:-1, 1:-1] += c

# treat the special cases of the borders
a[0, 1:-1] += .2 * b[0, :]
a[-1, 1:-1] += .2 * b[-1, :]
a[1:-1, 0] += .2 * b[:, 0]
a[1:-1, -1] += .2 * b[:, -1]

```

It gives the following result, which is the same as you nested loops.

```
[[  0.    2.2   3.4   4.6   4. ]
 [  6.2   2.6   4.2   3.   10.6]
 [  0.    3.4   4.8   6.2   4. ]
 [  6.2   2.6   4.2   3.   10.6]
 [  0.    2.2   3.4   4.6   4. ]]

```
I came up with this solution:

```
a = np.array([[0,0,0,0,0],
              [0,6,2,8,0],
              [0,1,5,3,0],
              [0,6,7,8,0],
              [0,0,0,0,0]],dtype= np.float32)

up= np.zeros_like(a)
down= np.zeros_like(a)
right= np.zeros_like(a)
left = np.zeros_like(a)

def new_a(a,u,r,d,l):

    c = np.copy(a)
    c[c <= 3] = 0

    up[:-2, 1:-1] += c[1:-1,1:-1] / 5.
    down[2:, 1:-1] += c[1:-1,1:-1] / 5.
    left[1:-1, :-2] += c[1:-1,1:-1]/ 5.
    right[1:-1, 2:] += c[1:-1,1:-1] / 5.
    a[1:-1,1:-1] -= c[1:-1,1:-1] * 4. / 5.
    a += up + down + left + right

    return a

```