如何在Python中迭代或递归识别二维数组的非nan相邻元素并完成全量标记

问题说明 #

原有代码仅单次遍历初始标记点的邻居，无法迭代扩展到所有连通区域，本质需要实现8邻接洪水填充（Flood Fill） 算法，在外层增加循环查找剩余未标记区域的最大值即可完成需求。

完整可运行代码 #

import numpy as np
from collections import deque

# 输入测试数组
test_arr = np.array([[1,2,np.nan,np.nan,5], 
                     [np.nan,2,np.nan,5,np.nan],
                     [np.nan,2,4,np.nan,6],
                     [np.nan,np.nan,np.nan,5,5],
                     [np.nan,np.nan,np.nan,np.nan,np.nan], 
                     [1,2,4,np.nan,np.nan],
                     [1,2,np.nan,np.nan,5],
                     [np.nan,4,np.nan,np.nan,6]])

row, col = test_arr.shape
# 初始化标记数组
res = np.zeros_like(test_arr, dtype=int)
# 8邻接方向定义
dirs = [(-1,-1), (-1,0), (-1,1),
        (0,-1),          (0,1),
        (1,-1),  (1,0),  (1,1)]

while True:
    # 筛选所有未标记的非nan元素坐标
    unmarked_coords = np.where((res == 0) & (~np.isnan(test_arr)))
    if len(unmarked_coords[0]) == 0:
        # 所有非nan元素已标记，结束循环
        break
    # 查找未标记元素的最大值
    current_max = np.nanmax(test_arr[unmarked_coords])
    # 取第一个最大值的坐标作为填充起点
    start_x, start_y = np.argwhere((test_arr == current_max) & (res == 0))[0]
    # BFS实现8连通洪水填充
    queue = deque()
    queue.append((start_x, start_y))
    res[start_x][start_y] = 1
    while queue:
        x, y = queue.popleft()
        for dx, dy in dirs:
            nx, ny = x + dx, y + dy
            # 边界校验 + 未标记校验 + 非nan校验
            if 0 <= nx < row and 0 <= ny < col and res[nx][ny] == 0 and not np.isnan(test_arr[nx][ny]):
                res[nx][ny] = 1
                queue.append((nx, ny))

print(res)

运行结果 #

[[1 1 0 0 1]
 [0 1 0 1 0]
 [0 1 1 0 1]
 [0 0 0 1 1]
 [0 0 0 0 0]
 [1 1 1 0 0]
 [1 1 0 0 1]
 [0 1 0 0 1]]

和需求预期的最终结果完全一致。

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