如何用纯NumPy实现二维数组分块平均值填充（无需循环）

问题描述 #

给定二维NumPy数组（各维度均可被指定的proximity值整除），需要按proximity大小分块，将每个块内的所有元素替换为该块的平均值。现有双重循环实现方案，寻求纯NumPy无循环的向量化解决方案。

示例输入

import numpy as np

my_array = np.array([
    [1, 2, 3, 4],
    [1, 2, 3, 4],
    [2, 3, 4, 5],
    [2, 3, 4, 5]
])
proximity = 2

示例输出

array([[1.5, 1.5, 3.5, 3.5],
       [1.5, 1.5, 3.5, 3.5],
       [2.5, 2.5, 4.5, 4.5],
       [2.5, 2.5, 4.5, 4.5]])

现有循环实现

proximity = 2
for i in range(0, my_array.shape[0], proximity):
    for j in range(0, my_array.shape[1], proximity):
        ii = i + proximity 
        jj = j + proximity 
        my_array[i:ii, j:jj] = np.average(my_array[i:ii, j:jj])

无循环解决方案 #

利用NumPy的形状重塑和向量化运算实现，核心思路是先将数组按块分组，计算块平均值后再扩展回原形状：

import numpy as np

# 初始化输入数组
my_array = np.array([
    [1, 2, 3, 4],
    [1, 2, 3, 4],
    [2, 3, 4, 5],
    [2, 3, 4, 5]
])
proximity = 2

# 1. 重塑数组，将每个proximity×proximity的块单独分组
# 新形状格式：(块行数, 块内行数, 块列数, 块内列数)
block_rows = my_array.shape[0] // proximity
block_cols = my_array.shape[1] // proximity
blocked_arr = my_array.reshape(block_rows, proximity, block_cols, proximity)

# 2. 计算每个块的平均值，保留维度以便后续扩展
block_means = blocked_arr.mean(axis=(1, 3), keepdims=True)

# 3. 将平均值重复填充回原块大小，再重塑回原数组形状
result = block_means.repeat(proximity, axis=1).repeat(proximity, axis=3).reshape(my_array.shape)

print(result)

步骤解释 #

1. 形状重塑：把原数组从(M, N)拆分为(M/prox, prox, N/prox, prox)，这样每个prox×prox的块成为独立的子数组，方便批量计算平均值。
2. 计算块平均值：通过mean方法指定axis=(1,3)，对每个块的行和列维度求平均，keepdims=True保持结果的维度与重塑后的数组匹配，避免后续扩展时维度不兼容。
3. 扩展平均值：用repeat方法将每个块的平均值重复proximity次，填充回原块的大小，最后重塑回原数组的形状。

优势 #

相比循环实现，该方案完全利用NumPy的向量化操作，避免了Python循环的开销，在处理大数组时效率提升明显。

内容的提问来源于stack exchange，提问作者wallstalk