python - 如何在 Python 中添加二维数组的相邻元素而不必使用嵌套循环？

Question

我想添加数组的相邻元素“3x3”并创建新数组。使用嵌套循环时，这段代码将被调用数千次，因此需要时间。

import tensorflow as tf
import numpy as np

rows=6
cols=8

array1 = np.random.randint(10, size=(rows, cols))
print(array1)

array2=np.zeros((rows-2,cols-2))

for i in range(rows-2):
    for j in range(cols-2):
        array2[i,j]=np.sum(array1[i:i+3,j:j+3])# print()
        
print("output\n",array2)


##My output
    [[9 4 9 6 1 4 9 0]
     [2 3 4 2 0 0 9 0]
     [2 8 9 7 6 9 4 8]
     [6 3 6 7 7 0 7 5]
     [2 1 4 1 7 6 9 9]
     [1 1 2 6 3 8 1 4]]
    output
     [[50. 52. 44. 35. 42. 43.]
     [43. 49. 48. 38. 42. 42.]
     [41. 46. 54. 50. 55. 57.]
     [26. 31. 43. 45. 48. 49.]]

通过矢量化，这可以解决。然而，我尝试了不同的技术，但从未有过任何运气，例如 reshape 然后添加数组，只使用一个带有大小行或列的循环。注意:在我的项目中，行和列的大小可能非常大。它类似于带有内核的 2D 卷积。问题是，是否可以在不使用循环的情况下实现这一点？，或者至少将其降低为具有较小的时间复杂度“仅采用行或列作为循环的大小”。

Answer 1

@trenixjetix 的微基准 convolve2d针对@n-ata 的 partial_sum 的 numpy 实现的解决方案.事实证明，即使对于相对较小的内核，第二种方法也更快。

import numpy as np
from scipy.signal import convolve2d

def partial_sum(array1, n):
    array1 = array1.cumsum(1).cumsum(0)
    res = array1.copy()
    res[:,n:] -= array1[:,:-n]
    res[n:] -= array1[:-n]
    res[n:,n:] += array1[:-n,:-n]
    return res[n-1:,n-1:]

rows = 500
cols = 500
n = 50   # relatively small n, since convolve2d becomes slow
np.random.seed(42)
array1 = np.random.randint(10, size=(rows, cols))

%timeit convolve2d(array1, np.ones((n,n), int), 'valid')
#1 loop, best of 5: 1.3 s per loop

%timeit partial_sum(array1, n)
#100 loops, best of 5: 3.36 ms per loop

partial_sum 的正确性通过与 convolve2d 结果的比较来验证。

np.testing.assert_equal(partial_sum(array1, n), convolve2d(array1, np.ones((n,n), int), 'valid'))

---

增加 n

的运行时复杂度

用于此基准测试的代码

import perfplot

perfplot.show(
    setup=lambda n: (np.random.randint(10, size=(500, 500)), n),
    kernels=[
        lambda a, n: partial_sum(a, n),
        lambda a, n: convolve2d(a, np.ones((n,n), int), 'valid')
    ],
    labels=['partial_sum','convolve2d'],
    n_range=[k for k in range(10,80,5)]
)