如何优化点邻域计算循环的执行效率？

问题背景 #

有一个名为zone的DataFrame，包含整数类型的x和y列代表点坐标，需要计算每个点的一阶邻居（切比雪夫距离=1，即周围3x3范围内除自身外的点）和二阶邻居（切比雪夫距离=2，即5x5范围内除自身和一阶邻居外的点）。原使用iterrows遍历的Pandas代码和左连接分组的SQL代码效率极低，百万级数据无法在合理时间内完成计算。

Pandas优化方案 #

方案1：基于KDTree的快速空间查询 #

利用scipy.spatial.KDTree实现O(n log n)复杂度的邻居查询，切比雪夫距离可直接通过指定metric='chebyshev'实现。

import numpy as np
import pandas as pd
from scipy.spatial import KDTree

# 生成测试数据（模拟用户场景）
data = np.random.randint(1000, 6000, size=(600000, 2))
zone = pd.DataFrame(data, columns=['x', 'y']).drop_duplicates().reset_index(drop=True)

# 构建KDTree，使用切比雪夫距离度量
tree = KDTree(zone[['x', 'y']], metric='chebyshev')

# 查询每个点半径≤1的邻居数（包含自身，需减1）
num_1st_neigh = tree.query_ball_point(zone[['x', 'y']], r=1, return_length=True) - 1

# 查询半径≤2的邻居数，减去自身和一阶邻居数得到二阶邻居数
total_r2 = tree.query_ball_point(zone[['x', 'y']], r=2, return_length=True) - 1
num_2nd_neigh = total_r2 - num_1st_neigh

# 合并结果到原DataFrame
zone['num_1st_neigh'] = num_1st_neigh
zone['num_2nd_neigh'] = num_2nd_neigh

优势：KDTree的空间查询效率远高于逐行遍历，百万级数据可在数秒内完成计算。

方案2：集合查找+偏移量匹配 #

将坐标转换为元组集合实现O(1)时间复杂度的存在性检查，预定义一阶/二阶所有偏移点，直接统计存在的偏移点数量。

import numpy as np
import pandas as pd

data = np.random.randint(1000, 6000, size=(600000, 2))
zone = pd.DataFrame(data, columns=['x', 'y']).drop_duplicates().reset_index(drop=True)

# 把坐标转成元组集合，用于快速查找
coord_set = set(zip(zone['x'], zone['y']))

# 定义一阶邻居的所有偏移（排除自身）
first_order_offsets = [(dx, dy) for dx in (-1,0,1) for dy in (-1,0,1) if (dx, dy) != (0,0)]
# 定义二阶邻居的所有偏移（仅切比雪夫距离=2的点）
second_order_offsets = [(dx, dy) for dx in (-2,-1,0,1,2) for dy in (-2,-1,0,1,2) if max(abs(dx), abs(dy)) == 2]

# 统计每个点的一阶邻居数
def count_first_neighbors(row):
    x, y = row['x'], row['y']
    return sum(1 for dx, dy in first_order_offsets if (x+dx, y+dy) in coord_set)

# 统计每个点的二阶邻居数
def count_second_neighbors(row):
    x, y = row['x'], row['y']
    return sum(1 for dx, dy in second_order_offsets if (x+dx, y+dy) in coord_set)

# 向量化计算
zone['num_1st_neigh'] = zone.apply(count_first_neighbors, axis=1)
zone['num_2nd_neigh'] = zone.apply(count_second_neighbors, axis=1)

优势：无需额外依赖，逻辑直观，百万级数据处理时间在分钟级以内。

SQL优化方案 #

方案1：偏移量预定义+EXISTS精确匹配 #

避免大表笛卡尔积，通过预定义所有一阶/二阶偏移量，用EXISTS检查偏移后的坐标是否存在，减少连接数据量。

WITH coord_offsets AS (
    -- 一阶邻居偏移（切比雪夫距离=1）
    SELECT dx, dy, 1 AS neighbor_level 
    FROM (VALUES (-1,-1), (-1,0), (-1,1), (0,-1), (0,1), (1,-1), (1,0), (1,1)) AS t(dx, dy)
    UNION ALL
    -- 二阶邻居偏移（切比雪夫距离=2）
    SELECT dx, dy, 2 AS neighbor_level 
    FROM (VALUES (-2,-2), (-2,-1), (-2,0), (-2,1), (-2,2),
                 (-1,-2), (-1,2),
                 (0,-2), (0,2),
                 (1,-2), (1,2),
                 (2,-2), (2,-1), (2,0), (2,1), (2,2)) AS t(dx, dy)
)
SELECT 
    t0.x,
    t0.y,
    SUM(CASE WHEN co.neighbor_level = 1 THEN 1 ELSE 0 END) AS num_1_neighbors,
    SUM(CASE WHEN co.neighbor_level = 2 THEN 1 ELSE 0 END) AS num_2_neighbors
FROM "table" t0
LEFT JOIN coord_offsets co 
    ON EXISTS (
        SELECT 1 FROM "table" t1 
        WHERE t1.x = t0.x + co.dx 
          AND t1.y = t0.y + co.dy
    )
GROUP BY t0.x, t0.y;

优化点：

1. 预定义偏移量避免了大范围的范围查询
2. EXISTS仅检查存在性，比JOIN生成更少的中间数据
3. 若x和y列有索引，EXISTS查询会极快

方案2：空间索引+距离筛选（适用于支持空间扩展的数据库） #

如果数据库支持空间索引（如PostgreSQL的PostGIS），将坐标转为空间点，利用空间索引快速筛选指定距离内的点。

-- 先创建空间索引（仅需执行一次）
CREATE INDEX idx_table_geom ON "table" USING GIST (ST_MakePoint(x, y));

SELECT 
    t0.x,
    t0.y,
    COUNT(CASE WHEN max(abs(t0.x - t1.x), abs(t0.y - t1.y)) = 1 THEN 1 END) AS num_1_neighbors,
    COUNT(CASE WHEN max(abs(t0.x - t1.x), abs(t0.y - t1.y)) = 2 THEN 1 END) AS num_2_neighbors
FROM "table" t0
LEFT JOIN "table" t1 
    -- ST_DWithin第三个参数为半径，TRUE表示使用L∞（切比雪夫）距离
    ON ST_DWithin(ST_MakePoint(t0.x, t0.y), ST_MakePoint(t1.x, t1.y), 2, TRUE)
    AND (t0.x != t1.x OR t0.y != t1.y)
GROUP BY t0.x, t0.y;

优势：空间索引能快速过滤出目标范围内的点，大幅减少JOIN的数据量，比原范围条件查询效率提升数倍。

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