python - 如何在 python 中向量化这个 for 循环？

Question

代码如下

import numpy as np


data = np.random.randint(0, 10, 12).reshape(3, 4)
print(data)

h, w = data.shape[:2]
dataMask = np.zeros((h, w, 10), np.int)

r = 2

for i in range(h):
    for j in range(w):
        for ir in range(i - r, i + r):
            for jr in range(j - r, j + r):
                if ir >= 0 and ir < h and jr >= 0 and jr < w:
                    dataMask[i, j, data[ir, jr]] += 1

print(dataMask)

我有一个形状为 (h, w) 的 numpy 数组“数据”。它的元素是整数 ∈[0, 10).
我创建了一个形状为 (h, w, 10) 的数组 dataMask。 dataMask[i, j, k]表示data中一个区域内值为k的点的个数。数据中的这个区域以(i,j)为中心，r = 2，是一个正方形。

如何向量化代码中的那些 for 循环？谢谢!

Answer 1

这是使用 cumsum 的一种方法:

import numpy as np


data = np.random.randint(0, 10, 1200).reshape(30, 40)
print(data)

h, w = data.shape[:2]
dataMask = np.zeros((h, w, 10), np.int)

r = 20

from time import time
T = []

T.append(time())

for i in range(h):
    for j in range(w):
        for ir in range(i - r, i + r):
            for jr in range(j - r, j + r):
                if ir >= 0 and ir < h and jr >= 0 and jr < w:
                    dataMask[i, j, data[ir, jr]] += 1

T.append(time())

m1 = np.zeros((h, w, 10), np.int)
np.put_along_axis(m1, data[...,None], 1, 2)
m2 = np.empty_like(m1)
m1 = m1.cumsum(1)
m2[: ,:-r+1] = m1[:, r-1:]
m2[:, -r+1:] = m1[:, -1, None]
m2[:, r+1:] -= m1[:, :-r-1]
m2 = m2.cumsum(0)
m1[:-r+1] = m2[r-1:]
m1[-r+1:] = m2[-1, None]
m1[r+1:] -= m2[:-r-1]

T.append(time())


assert (dataMask==m1).all()

print(np.diff(T))

使用 h,w,r = 30,40,20 运行示例

# time [seconds] used by
# OP            cumsum
[9.23162699e-01 3.41892242e-04]