优化CSV文件中虚假树木坐标识别代码的方案探讨

针对5万条树木数据的虚假识别需求，原双重循环O(n²)的复杂度完全无法支撑，以下是几种高效的优化思路及实现：

1. 空间索引（KD-Tree/球树）快速定位邻近树木 #

核心是避免遍历所有树木对，只筛选出可能满足条件的邻近目标：

• 逻辑转化：原判定条件 d = e - (r1 + r2) < 0.5 等价于 e < r1 + r2 + 0.5。对任意树木i，只需搜索以其坐标为中心、半径为 r_i + max_r + 0.5 的范围内的树木（max_r为所有树木半径的最大值），再在这些候选中验证条件。
• 实现代码（用scipy的cKDTree，C实现比纯Python快数倍）：

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

# 读取数据并预处理
df = pd.read_csv('trees.csv')
df['r'] = df['直径'] / 2
coords = df[['x', 'y']].values
radii = df['r'].values
max_r = radii.max()

# 构建KD-Tree索引
tree = cKDTree(coords)

# 获取所有符合条件的树木对
search_radii = radii + max_r + 0.5
# 生成所有i<j的候选对，避免重复计算
candidate_pairs = tree.query_pairs(search_radii, output_type='ndarray')

# 向量化计算验证条件
distances = np.linalg.norm(coords[candidate_pairs[:,0]] - coords[candidate_pairs[:,1]], axis=1)
d_values = distances - (radii[candidate_pairs[:,0]] + radii[candidate_pairs[:,1]])
valid_pairs = candidate_pairs[d_values < 0.5]

# 标记虚假树木
fake_ids = np.unique(valid_pairs.flatten())
df.loc[fake_ids, 'is_fake'] = True

# 保存结果
df.to_csv('trees_with_fake_marked.csv', index=False)

• 复杂度：O(n log n)，5万条数据可在数秒内处理完成。

2. 空间网格划分（轻量无依赖方案） #

如果不想依赖第三方库，可手动实现网格索引：

• 步骤：
  1. 设定网格大小为 max_r + 0.5，确保只有同一/相邻网格的树木才可能满足条件。
  2. 给每棵树分配网格坐标：grid_x = int(x // grid_size)，grid_y = int(y // grid_size)。
  3. 对每个网格，仅与自身及8个相邻网格内的树木计算距离并验证条件。
• 核心代码片段：

grid_size = max_r + 0.5
# 构建网格到树木索引的映射
grid_map = {}
for idx, (x, y) in enumerate(coords):
    gx, gy = int(x // grid_size), int(y // grid_size)
    grid_map.setdefault((gx, gy), []).append(idx)

fake_ids = set()
for idx in range(len(df)):
    x, y = coords[idx]
    gx, gy = int(x // grid_size), int(y // grid_size)
    # 检查当前网格及8个相邻网格
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            neighbor_grid = (gx + dx, gy + dy)
            if neighbor_grid not in grid_map:
                continue
            for j in grid_map[neighbor_grid]:
                if idx == j:
                    continue
                e = np.linalg.norm(coords[idx] - coords[j])
                d = e - (radii[idx] + radii[j])
                if d < 0.5:
                    fake_ids.add(idx)
                    fake_ids.add(j)
                    break
        if idx in fake_ids:
            break

df['is_fake'] = df.index.isin(fake_ids)

3. 结果优化建议 #

• 原逻辑中满足条件的两棵树会被同时标记为虚假，可根据业务需求调整：比如仅标记直径更小的树木，或保留其中一棵，减少不必要的误标。
• 若内存不足，可分批次处理数据：按x或y坐标切分数据集，处理完一批后再处理下一批，最后合并结果。

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