二维二进制数组中各像素到最近0值的距离计算优化问询

给定仅包含0和1的二维二进制数组，需将所有值为1的像素替换为其到最近0值像素的距离（对角距离计为1，即采用切比雪夫距离）。示例如下：

输入：
[[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 0, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
]
输出：
[[2, 2, 2, 2, 2],
[2, 1, 1, 1, 2],
[2, 1, 0, 1, 2],
[2, 1, 1, 1, 2],
[2, 2, 2, 2, 2],
]

现有暴力解法的问题 #

当前暴力解法通过枚举所有0和1的像素索引，计算两两之间的欧氏距离矩阵并取最小值，但该方法时间复杂度为O((n*m)^2)，空间复杂度也极高，处理2k×2k的大矩阵时会直接出现内存溢出、运算速度极慢的问题。暴力解法代码如下：

import numpy as np
# example sample
maze = np.ones((100, 100))
maze[40:60, 40:60] = 0

# find indices
mask_z = maze == 0
idx_z = np.nonzero(mask_z)
idx_nz = np.nonzero(~mask_z)
idx_z = np.stack(idx_z, axis=-1)
idx_nz = np.stack(idx_nz, axis=-1)

# calculate distance matrix
dist_matrix = np.linalg.norm(idx_nz.reshape(-1, 1, 2) - idx_z, axis=-1)
dists = np.min(dist_matrix, axis=-1)

# assign distances
mapped = np.zeros(maze.shape[:2], dtype=int)
mapped[tuple(idx_nz.T)] = dists.squeeze()

优化方案：多源广度优先搜索（BFS） #

针对这类多源最短路径问题，多源BFS是最优解决方案。该方法从所有0像素同时出发，逐层向外扩展，每个1像素第一次被访问时的层数即为其到最近0的距离，时间复杂度为O(nm)，空间复杂度为O(nm)，完全适配大矩阵场景。

代码实现 #

import numpy as np
from collections import deque

def nearest_zero_distance(maze):
    rows, cols = maze.shape
    # 初始化距离矩阵：0的位置设为0，1的位置设为-1（标记未访问）
    dist = np.full((rows, cols), -1, dtype=int)
    q = deque()
    
    # 批量加入所有0像素的坐标到队列
    zero_indices = np.argwhere(maze == 0)
    for x, y in zero_indices:
        dist[x][y] = 0
        q.append((x, y))
    
    # 定义8个移动方向（包含对角线，对应切比雪夫距离）
    directions = [(-1,-1), (-1,0), (-1,1),
                  (0,-1),          (0,1),
                  (1,-1),  (1,0), (1,1)]
    
    # 执行多源BFS
    while q:
        x, y = q.popleft()
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            # 检查坐标是否在矩阵范围内，且未被访问过
            if 0 <= nx < rows and 0 <= ny < cols and dist[nx][ny] == -1:
                dist[nx][ny] = dist[x][y] + 1
                q.append((nx, ny))
    return dist

# 测试示例输入
test_maze = np.ones((5,5))
test_maze[2][2] = 0
print(nearest_zero_distance(test_maze))

# 处理2k×2k大矩阵示例
large_maze = np.ones((2000, 2000))
large_maze[800:1200, 800:1200] = 0
large_dist = nearest_zero_distance(large_maze)

方案优势 #

• 时间高效：每个像素仅被访问一次，处理2k×2k矩阵仅需约400万次操作，远优于暴力解法的平方级复杂度
• 内存可控：无需存储巨大的距离矩阵，仅需维护队列和距离矩阵，内存占用与矩阵大小线性相关
• 结果准确：严格按照切比雪夫距离计算最近0的距离，完全符合需求

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