NumPy数组如何替换上方、左侧相邻元素均为5的数值4

实现思路 #

不需要逐元素循环遍历，直接利用NumPy的切片偏移特性做矢量化判断即可，性能远高于手写循环，同时自动规避数组边界越界问题：

• 目标替换元素必须同时满足3个条件：
  • 元素自身取值为4
  • 元素正上方相邻位置取值为5
  • 元素正左侧相邻位置取值为5
• 由于第一行元素无上方邻居、第一列元素无左侧邻居，天然不满足替换条件，直接从第2行第2列（0索引下的[1:,1:]切片）开始判断即可，不会出现索引越界。

完整实现代码 #

import numpy as np

# 1. 处理输入为NumPy整数数组（如果已经是数值数组可跳过这步）
input_str = """0000300 
0005000
0054000
0000045
0002050"""
arr = np.array(
    [[int(c) for c in line.strip()] for line in input_str.splitlines()],
    dtype=int
)
result = arr.copy()  # 拷贝数组避免修改原数据

# 2. 生成替换掩码：三个条件同时为True的位置才需要替换
# 切片对齐逻辑：
# result[1:, 1:] 对应所有可能有上、左邻居的元素位置
# result[:-1, 1:] 对应上述位置的正上方邻居
# result[1:, :-1] 对应上述位置的正左侧邻居
replace_mask = (
    (result[1:, 1:] == 4)
    & (result[:-1, 1:] == 5)
    & (result[1:, :-1] == 5)
)

# 3. 执行替换
result[1:, 1:][replace_mask] = 8

# 如果需要转回和输入一致的每行数字串格式，用下面代码
output_str = '\n'.join(''.join(row.astype(str)) for row in result)
print(output_str)

示例运行说明 #

针对给出的测试用例，代码会正确定位到第三行第四列（0索引位置(2,3)）的4：该位置上方是第二行第四列的5，左侧是第三行第三列的5，满足替换条件，会被替换为8，其余元素保持原值。

注：给出的预期输出最后一行为0002000，和「其余元素保持原值不变」的规则冲突（原输入最后一行第6位是5，无符合条件的替换触发逻辑），属于示例笔误，代码运行后对应位置会保留原值5，最后一行输出为0002050。

内容的提问来源于stack exchange，提问作者Loïc