用numpy避免for循环加快计算速度

Question

我正在使用1800 x 900矩阵。对于矩阵中的每个单元格，我需要做一个类似的操作，可以用一个例子来描述：

a1b1 a1b2 a1b3 .... a1b900
a2b1 a2b2 ..
.
a1800b1 ... ...   ..a1800b900

对于一个新矩阵中的每个单元格，我希望它的值等于它的左上、右上和右上邻居的和。所以，我想要a2b2 = a1b1 + a1b2 + a1b3。对于顶行，它可以是与当前单元格相同的值。你知道吗

我可以很容易地使用2 for循环来实现这一点，但是有没有办法将其矢量化，从而加快整个过程？你知道吗

Answer 1

你可以像下面那样把它矢量化

arr = np.arange(16).reshape(4,4)

arr2 = np.pad(arr, [(1,0), (1,1)], mode="constant")[:-1, :]

result = arr+arr2[:,0:-2]+arr2[:,1:-1]+arr2[:,2:]

描述

• np.pad(arr, [(1,0), (1,1)], mode="constant")在矩阵的左、右和顶部填充零。索引[：-1，：]将删除最后一行，有效地将矩阵移动一行
• arr2[:,0:-2], arr2[:,1:-1], arr2[:,2:]将为arr的每个对应元素提供左上角、右上角和右上角元素。如果没有元素，则给出0（上一步中填充的零）

结果

## arr
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

## result
array([[ 0,  1,  2,  3],
       [ 5,  8, 12, 12],
       [17, 24, 28, 24],
       [29, 40, 44, 36]])