如何基于K近邻计算Z值平均值？求可扩展Python函数

修正原代码的近邻拟合逻辑 #

首先需要注意：你注释中提到“nbrs reads in the x,y values only”，但原代码直接用了包含Z值的整个X数组拟合模型，这会导致近邻计算错误。正确的做法是只提取x,y坐标用于近邻计算，Z值单独提取：

from sklearn.neighbors import NearestNeighbors
import numpy as np

X = np.array([[6,-3, 0.1], [-5,-9, 0.5], [3,-7, 0.8], [-10,6, 0.5], [-4,-16, 0.9], [1,-0.5, 0]])
# 分离x,y坐标和Z值
xy_coords = X[:, :2]
z_values = X[:, 2]

# 初始化K近邻模型，K可设置为任意值（如2、50）
K = 2
nbrs = NearestNeighbors(n_neighbors=K).fit(xy_coords)
distances, indices = nbrs.kneighbors(xy_coords)

通用平均Z值计算函数 #

下面的函数支持任意K值，利用numpy向量化操作实现高效计算：

def compute_avg_z(indices, z_values):
    # 通过索引获取每个点的所有近邻Z值，按行计算平均值
    neighbor_z_values = z_values[indices]
    return neighbor_z_values.mean(axis=1)

调用示例

average_z = compute_avg_z(indices, z_values)
print(average_z)

可选：排除当前点自身的情况 #

由于NearestNeighbors默认会将点自身作为第一个近邻，如果需要排除自身计算平均，只需调整索引范围：

def compute_avg_z_exclude_self(indices, z_values):
    # 跳过第一个索引（自身），取其余近邻的Z值计算平均
    neighbor_z_values = z_values[indices[:, 1:]]
    return neighbor_z_values.mean(axis=1)

说明 #

• 该实现依赖numpy的向量化运算，无论K=2还是K=50，都能高效处理大规模数据，比循环遍历的方式性能更优。
• 函数的参数indices是NearestNeighbors返回的近邻索引数组，z_values是所有点的Z值数组，两者维度匹配即可正常运行。

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