如何按特定条件修改NumPy数组元素：保留3×3及以上连续1

问题描述 #

给定如下NumPy数组：

import numpy as np

data = np.array([[4,4,4,0,1,1,1,0,0,0,0,1,0,0,1],
           [3,0,0,1,1,1,1,1,1,1,1,0,0,1,0],
           [6,0,0,1,1,1,1,1,0,0,0,0,1,0,0]])

需求：仅保留数组中由1构成的3×3及更大尺寸的连续正方形区域，其余所有不符合要求的1替换为0，非1元素保持不变。

预期输出：

[[4 4 4 0 1 1 1 0 0 0 0 0 0 0 0]
 [3 0 0 0 1 1 1 0 0 0 0 0 0 0 0]
 [6 0 0 0 1 1 1 0 0 0 0 0 0 0 0]]

解决方案 #

通过滑动窗口检测全1的3×3区域，标记需要保留的1位置，最后生成结果：

import numpy as np

# 原输入数组
data = np.array([[4,4,4,0,1,1,1,0,0,0,0,1,0,0,1],
           [3,0,0,1,1,1,1,1,1,1,1,0,0,1,0],
           [6,0,0,1,1,1,1,1,0,0,0,0,1,0,0]])

# 生成布尔掩码：标记所有值为1的位置
mask = data == 1
rows, cols = mask.shape

# 创建保留标记数组，初始全为False
keep = np.zeros_like(mask, dtype=bool)

# 遍历所有可容纳3×3窗口的左上角坐标
for i in range(rows - 2):
    for j in range(cols - 2):
        # 检查当前3×3窗口内是否全为1
        if np.all(mask[i:i+3, j:j+3]):
            # 标记该窗口内的所有位置为需保留
            keep[i:i+3, j:j+3] = True

# 生成最终结果：保留符合要求的1，其余1置0，非1元素不变
result = np.where(keep, data, np.where(mask, 0, data))

print(result)

代码说明 #

1. 掩码生成：用data == 1提取所有1的位置，得到布尔掩码mask，简化后续区域判断。
2. 窗口检测：遍历所有能放下3×3窗口的左上角位置，检查窗口内是否全为1。若满足条件，将该窗口内的所有位置标记为需保留。
3. 结果生成：通过np.where嵌套判断，对每个元素做处理：
   • 若属于需保留的1区域，保留原数值；
   • 若为1但不属于保留区域，替换为0；
   • 非1元素直接保留原值。

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