如何在非空间操作下从网格DataFrame中查找种子点的K近邻？

需求说明 #

• 现有两个DataFrame：
  • grid：存储方形网格节点，包含字符串类型id、整数类型x/y坐标，部分节点缺失
  • seeds：存储网格中的种子节点，包含字符串类型seed_id、整数类型x/y坐标
• 需为每个seed_id筛选出前K个近邻节点，规则为：
  • 一级近邻：种子点[x-1, x+1]×[y-1, y+1]范围内的所有节点（无缺失时共8个）
  • 二级近邻：种子点[x-2, x+2]×[y-2, y+2]范围内的所有节点（叠加一级后无缺失时共24个）
  • 以此类推，按层级从近到远取节点，直到凑够K个

实现思路 #

利用网格坐标为整数的特性，通过纯Pandas操作完成，无需空间几何库：

1. 计算覆盖K个节点所需的最大扩展层级n（公式推导：每级新增8×n个节点，总节点数为4n(n+1)，找到满足4n(n+1)≥K的最小n）
2. 通过笛卡尔积关联seeds和grid，筛选出种子节点x±n、y±n范围内的节点，排除种子自身
3. 按节点与种子的层级（最大坐标偏移量）、坐标排序，分组取前K个节点

代码实现 #

import pandas as pd

def get_k_neighbors(grid, seeds, K):
    # 计算需要扩展的最大层级n
    n = 0
    while 4 * n * (n + 1) < K:
        n += 1

    # 笛卡尔积关联种子和网格节点
    seeds['join_key'] = 1
    grid['join_key'] = 1
    cross_join = seeds.merge(grid, on='join_key').drop('join_key', axis=1)

    # 筛选范围内节点，排除种子自身
    valid_mask = (
        (cross_join['x_x'] - n <= cross_join['x_y']) & (cross_join['x_y'] <= cross_join['x_x'] + n) &
        (cross_join['y_x'] - n <= cross_join['y_y']) & (cross_join['y_y'] <= cross_join['y_x'] + n) &
        ~((cross_join['x_x'] == cross_join['x_y']) & (cross_join['y_x'] == cross_join['y_y']))
    )
    candidates = cross_join[valid_mask]

    # 计算层级（最大坐标偏移量）用于排序
    candidates['level'] = candidates.apply(
        lambda row: max(abs(row['x_x'] - row['x_y']), abs(row['y_x'] - row['y_y'])),
        axis=1
    )

    # 按种子分组，先按层级由近到远，再按坐标排序，取前K个节点
    result = candidates.groupby('seed_id').apply(
        lambda group: group.sort_values(['level', 'x_y', 'y_y']).head(K)
    ).reset_index(drop=True)

    # 整理输出格式
    return result[['seed_id', 'id']]

示例验证 #

构造测试数据 #

# 网格节点数据（对应示例中的缺失网格）
grid_data = [
    ('1', 0, 3), ('2', 1, 3), ('3', 2, 3), ('4', 3, 3),
    ('4', 0, 2), ('5', 1, 2), ('6', 2, 2),
    ('7', 0, 1), ('8', 1, 1), ('9', 2, 1), ('10', 3, 1),
    ('11', 0, 0), ('12', 1, 0), ('13', 2, 0), ('14', 3, 0)
]
grid = pd.DataFrame(grid_data, columns=['id', 'x', 'y'])

# 种子节点数据
seeds_data = [('7', 0, 1), ('10', 3, 1)]
seeds = pd.DataFrame(seeds_data, columns=['seed_id', 'x', 'y'])

调用函数获取一级近邻（K=5） #

neighbors = get_k_neighbors(grid, seeds, K=5)
print(neighbors.to_csv(sep=',', index=False))

输出结果 #

seed_id,id
7,4
7,5
7,8
7,11
7,12
10,6
10,9
10,13
10,14

与需求示例完全一致。

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