基于NumPy计算三维网格点在14个方向上的最近0值邻居距离

核心思路 #

针对问题需求，我们可以通过轴向向量化计算+对角线分组处理的方式替代三重嵌套循环，最大化利用NumPy的向量化特性提升效率：

• 轴向（±x/±y/±z）直接用np.minimum.accumulate做累积计算，完全避免循环；
• 面对角线（8个方向）通过射线分组+二分查找的方式，大幅减少循环次数。

实现步骤与代码 #

1. 定义14个方向向量 #

import numpy as np

# 14个目标方向：6个轴向 + 8个面对角线
directions = [
    # 轴向
    (1, 0, 0), (-1, 0, 0),
    (0, 1, 0), (0, -1, 0),
    (0, 0, 1), (0, 0, -1),
    # 面对角线（xy/xz平面）
    (1, 1, 0), (1, -1, 0), (-1, 1, 0), (-1, -1, 0),
    (1, 0, 1), (1, 0, -1), (-1, 0, 1), (-1, 0, -1),
]

2. 主函数实现 #

def directional_nearest_zero(arr):
    nx, ny, nz = arr.shape
    output = np.zeros((nx, ny, nz, len(directions)), dtype=np.int32)
    mask_0 = arr == 0  # 0值点掩码

    for dir_idx, (di, dj, dk) in enumerate(directions):
        # ---------------------- 处理轴向方向 ----------------------
        if di != 0 and dj == 0 and dk == 0:
            idx = np.arange(nx)[:, None, None]
            zero_pos = np.where(arr == 0, idx, np.inf)
            if di == 1:
                # x正方向：从右往左累积最近0的位置
                nearest_zero = np.minimum.accumulate(zero_pos[::-1, :, :])[::-1, :, :]
                dist = nearest_zero - idx
            else:
                # x负方向：从左往右累积最近0的位置
                nearest_zero = np.minimum.accumulate(zero_pos, axis=0)
                dist = idx - nearest_zero
            dist[mask_0] = 0
            output[:, :, :, dir_idx] = dist
            continue

        elif dj != 0 and di == 0 and dk == 0:
            idx = np.arange(ny)[None, :, None]
            zero_pos = np.where(arr == 0, idx, np.inf)
            if dj == 1:
                nearest_zero = np.minimum.accumulate(zero_pos[:, ::-1, :])[:, ::-1, :]
                dist = nearest_zero - idx
            else:
                nearest_zero = np.minimum.accumulate(zero_pos, axis=1)
                dist = idx - nearest_zero
            dist[mask_0] = 0
            output[:, :, :, dir_idx] = dist
            continue

        elif dk != 0 and di == 0 and dj == 0:
            idx = np.arange(nz)[None, None, :]
            zero_pos = np.where(arr == 0, idx, np.inf)
            if dk == 1:
                nearest_zero = np.minimum.accumulate(zero_pos[:, :, ::-1])[:, :, ::-1]
                dist = nearest_zero - idx
            else:
                nearest_zero = np.minimum.accumulate(zero_pos, axis=2)
                dist = idx - nearest_zero
            dist[mask_0] = 0
            output[:, :, :, dir_idx] = dist
            continue

        # ---------------------- 处理面对角线方向 ----------------------
        # 构建射线标识与参数t：同一条射线上的点共享相同标识
        i, j, k = np.indices(arr.shape)
        if di != 0 and dj != 0:
            # xy平面对角线：射线标识为(i-j, k)，参数t用i表示
            ray_id = (i - j, k)
            t = i
        else:
            # xz平面对角线：射线标识为(i-k, j)，参数t用i表示
            ray_id = (i - k, j)
            t = i

        # 归一化射线标识为非负整数，方便分组
        norm_a = ray_id[0] - ray_id[0].min()
        norm_b = ray_id[1] - ray_id[1].min()
        max_b = norm_b.max() + 1
        group_id = norm_a * max_b + norm_b

        # 扁平化数据并按分组+参数t排序
        flat_t = t.flatten()
        flat_group = group_id.flatten()
        flat_mask0 = mask_0.flatten()
        flat_coords = np.stack([i.flatten(), j.flatten(), k.flatten()], axis=1)

        sort_idx = np.lexsort((flat_t, flat_group))
        sorted_t = flat_t[sort_idx]
        sorted_group = flat_group[sort_idx]
        sorted_mask0 = flat_mask0[sort_idx]
        sorted_coords = flat_coords[sort_idx]

        # 找到每个分组的边界
        group_boundaries = np.where(np.diff(sorted_group) != 0)[0] + 1
        group_boundaries = np.concatenate([[0], group_boundaries, [len(sorted_group)]])

        # 遍历每个分组计算距离
        for start, end in zip(group_boundaries[:-1], group_boundaries[1:]):
            group_t = sorted_t[start:end]
            group_mask0 = sorted_mask0[start:end]
            coords = sorted_coords[start:end]

            zero_positions = group_t[group_mask0]
            if len(zero_positions) == 0:
                continue  # 该射线无0值点，保持默认0

            # 二分查找找到每个点右侧最近的0值位置
            dists_idx = np.searchsorted(zero_positions, group_t, side='right')
            valid = dists_idx < len(zero_positions)
            group_dist = np.zeros_like(group_t)
            group_dist[valid] = zero_positions[dists_idx[valid]] - group_t[valid]
            group_dist[~valid] = np.inf  # 无0值点的位置设为无穷大，可按需修改

            # 赋值到输出数组
            output[coords[:, 0], coords[:, 1], coords[:, 2], dir_idx] = group_dist

        # 0值点强制设为0
        output[mask_0, dir_idx] = 0

    return output

效率优化关键点 #

1. 轴向全向量化：利用np.minimum.accumulate一次性完成整个维度的最近0值位置计算，无任何循环；
2. 对角线分组处理：将同方向射线的点分组，每组用np.searchsorted二分查找替代逐点遍历，循环次数从总点数降至射线组数；
3. 掩码批量处理：通过布尔掩码一次性将所有0值点的距离设为0，避免逐点判断。

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