如何在numpy二维数组中为每行查询非零值的k个最近邻点

我们可以按「提取有效非零点集→构建KNN索引→按行查询最近邻」的流程实现，默认基于点在二维数组中的坐标（行、列索引）的欧氏距离计算近邻，你可以根据需求调整距离规则。

步骤1：提取所有有效非零点 #

首先过滤出数组中所有数值大于0的点，同时记录它们的坐标和对应数值：

import numpy as np
from sklearn.neighbors import NearestNeighbors

# 你的原始二维数组
# data = ...

# 提取所有值>0的点的行、列索引
row_idx, col_idx = np.where(data > 0)
# 构造有效点坐标数组，每个元素格式为(行号, 列号)
valid_points = np.column_stack((row_idx, col_idx))
# 对应坐标的数值
valid_values = data[row_idx, col_idx]

步骤2：构建KNN查询索引 #

用有效点集训练KNN模型，支持自定义距离度量：

k = 5 # 可替换为你需要的近邻数量
# 构建KNN模型，metric参数可替换为'manhattan'(曼哈顿距离)、'cosine'(余弦距离)等
knn = NearestNeighbors(n_neighbors=k, metric='euclidean')
knn.fit(valid_points)

步骤3：按行查询k个最近邻 #

遍历每行，提取该行的非零点坐标，查询对应的k近邻：

# 存储结果的字典，key为行号，value为该行每个非零点对应的近邻信息
row_knn_result = {}

# 遍历所有行
for row in range(data.shape[0]):
    # 提取当前行的非零点坐标
    row_mask = row_idx == row
    if not np.any(row_mask):
        # 当前行无非零点，跳过
        continue
    current_points = valid_points[row_mask]
    # 查询k近邻，返回每个查询点对应的近邻距离、近邻在valid_points中的索引
    distances, nn_indices = knn.kneighbors(current_points)
    # 整理结果：每个点的原坐标、原数值、近邻坐标、近邻数值、近邻距离
    point_results = []
    for i in range(len(current_points)):
        point_info = {
            "current_coord": current_points[i],
            "current_value": data[tuple(current_points[i])],
            "nn_coords": valid_points[nn_indices[i]],
            "nn_values": valid_values[nn_indices[i]],
            "nn_distances": distances[i]
        }
        point_results.append(point_info)
    row_knn_result[row] = point_results

可选：纯Numpy实现（无需第三方库） #

如果不想依赖scikit-learn，可以直接计算距离矩阵后取前k小的结果，适合小数据集场景：

k = 5
# 计算所有有效点两两之间的欧氏距离矩阵
dist_matrix = np.linalg.norm(valid_points[:, None] - valid_points, axis=2)
# 每个点取距离最小的k个索引（需要排除自身的话可以取[:,1:k+1]）
nn_indices = np.argsort(dist_matrix, axis=1)[:, :k]
# 后续按行整理结果逻辑和上述一致

注意事项 #

• 如果有效点总数量小于k，需要提前做边界处理，避免查询报错
• 如果你的近邻计算规则是基于行向量的特征距离而非坐标距离，只需要把valid_points替换为排除0特征后的行向量即可
• 大数量级场景下建议用scikit-learn的KNN，底层用kd树/球树优化，查询效率远高于纯numpy暴力计算

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