基于NumPy将数组零值替换为参考数组最近非NaN值

需要将数组ar_rand中的所有0值，替换为数组ar_label中对应位置最近的非NaN值（基于二维欧氏距离计算）。要求优先使用NumPy或Xarray实现，且方案不依赖ar_label中非NaN值的特定分布。

示例数组定义 #

import numpy as np

nrow = 10
ar_label = np.arange(nrow**2).reshape(nrow, nrow)
ar_label[1:4, 1:4] = 100
ar_label[6:9, 2:5] = 200
ar_label[2:5, 6:9] = 300
ar_label = np.where(ar_label < 100, np.nan, ar_label)

print("ar_label:")
print(ar_label)

输出：

array([[ nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan],
       [ nan, 100., 100., 100.,  nan,  nan,  nan,  nan,  nan,  nan],
       [ nan, 100., 100., 100.,  nan,  nan, 300., 300., 300.,  nan],
       [ nan, 100., 100., 100.,  nan,  nan, 300., 300., 300.,  nan],
       [ nan,  nan,  nan,  nan,  nan,  nan, 300., 300., 300.,  nan],
       [ nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan],
       [ nan,  nan, 200., 200., 200.,  nan,  nan,  nan,  nan,  nan],
       [ nan,  nan, 200., 200., 200.,  nan,  nan,  nan,  nan,  nan],
       [ nan,  nan, 200., 200., 200.,  nan,  nan,  nan,  nan,  nan],
       [ nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan]])

np.random.seed(11)
ar_rand = np.random.randint(0, nrow*3, size=nrow**2).reshape(nrow, nrow)
ar_rand = np.where(ar_rand == 0, ar_rand, np.nan)

print("\nar_rand:")
print(ar_rand)

输出：

array([[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       [nan, nan, nan, nan, nan,  0., nan, nan, nan, nan],
       [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       [nan, nan, nan, nan, nan, nan, nan, nan,  0., nan],
       [nan,  0., nan, nan, nan, nan, nan, nan, nan, nan],
       [nan, nan, nan, nan, nan, nan,  0., nan, nan, nan],
       [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       [nan, nan, nan, nan, nan, nan,  0., nan, nan, nan],
       [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       [nan, nan,  0., nan, nan, nan, nan, nan, nan, nan]])

解决方案 #

基础NumPy实现 #

核心逻辑是提取非NaN点和零点的坐标，逐个计算欧氏距离并替换：

# 提取ar_label中非NaN值的坐标与对应值
label_coords = np.argwhere(~np.isnan(ar_label))
label_values = ar_label[~np.isnan(ar_label)]

# 提取ar_rand中0值的坐标
zero_coords = np.argwhere(ar_rand == 0)

# 遍历每个0点，替换为最近的非NaN值
for x, y in zero_coords:
    # 计算当前点到所有非NaN点的欧氏距离（省略平方根不影响大小比较，提升效率）
    distances = (x - label_coords[:, 0]) ** 2 + (y - label_coords[:, 1]) ** 2
    min_idx = np.argmin(distances)
    ar_rand[x, y] = label_values[min_idx]

print("替换后的ar_rand:")
print(ar_rand)

大规模数据优化方案（scipy加速） #

如果数组规模较大，用scipy.spatial.cKDTree可以大幅提升最近邻查找效率：

from scipy.spatial import cKDTree

# 构建非NaN点的KD树
tree = cKDTree(label_coords)
# 批量查询所有0点的最近邻索引
_, indices = tree.query(zero_coords)
# 批量替换0值
ar_rand[tuple(zero_coords.T)] = label_values[indices]

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