三维numpy数组零值元素填充：取最近非零元素（L1距离）

我需要将三维NumPy数组中的所有零值元素替换为最近的非零元素值，若存在多个等距的非零元素，任选其一即可。以下是二维场景下的效果示例：

示例1 #

输入：

0 0 0 0
0 1 2 0
0 3 4 0
0 0 0 0

输出：

1 1 2 2
1 1 2 2
3 3 4 4
3 3 4 4

示例2 #

输入：

1 0 2 0
0 3 0 1
2 0 4 0

有效输出：

1 2 2 2
3 3 2 1
2 3 4 1

我最初想到的方案是逐个处理非零索引，填充相邻零值元素，重复直至无零值，但该方法效率极低。由于要处理的数组规模约为256×256×256，急需高效解决方案。

现有可行代码（仍需优化） #

以下代码可行且速度尚可，但希望找到无循环的优化方案：

is_useful_value_data = np.isin(actual_data, values_we_want_to_preserve)
indices = np.where(is_useful_value_data)

interp = scipy.interpolate.NearestNDInterpolator(indices, actual_data[is_useful_value_data])

indices_to_fill = np.where(~is_useful_value_data)

for index in np.transpose(indices_to_fill):
    x = index[0]
    y = index[1]
    z = index[2]
    nearest_val = interp([x,y,z])
    actual_data[x,y,z] = int(nearest_val)

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