优化Python KD树搜索

Question

Scipy（http://www.scipy.org/）提供了两个KD树类：KDTree和cKDTree。

cKDTree速度更快，但可定制性和查询能力不如KDTree（从文档中可以看出）。

我的问题是：我有一个包含3百万个二维（X，Y）点的列表。我需要返回每个点距离X单位以内的所有点。

使用KDtree时，有一种选项可以做到这一点：KDtree.query_ball_tree()它生成一个列表，其中包含距离每个其他点X单位以内的所有点的列表。然而：此列表非常大，并且很快就会填满我的虚拟内存（约为7.44亿个项目长）。

解决方案1：是否有一种方法可以在写入时将此列表解析为文本文件？

解决方案2：我尝试使用for循环（对于列表中的每个点），然后通过使用KDtree.query_ball_point()查找该单个点在X单位内的邻居。然而：这需要运行数百万次查询，因此需要很长时间。是否有一个cKDTree等效于此KDTree工具？

解决方案3：我不知道，还有其他人有什么想法吗？

- Dlinet

Answer 1

从scipy 0.12开始，两个KD Tree类具有功能平等性。引用其公告：

cKDTree功能完备

KDTree的Cython版本cKDTree现在具备了完整的功能。大多数操作(构建、查询、query_ball_point、query_pairs、count_neighbors和sparse_distance_matrix)在cKDTree中比在KDTree中快200到1000倍。除极少数例外外，cKDTree与KDTree具有完全相同的接口，并且可以作为一个即插即用的替代品。

- jorgeca

---

啊，那太好了。我没有编译源代码的技能/经验，所以也许我会研究一下。否则，除非发布另一个解决方案，否则我将等待scipy的新版本发布。 - Dlinet

@Dlinet 0.12 版本上个月发布。 - jorgeca

对于任何在现代时代（2022年）查看此答案的人来说，从SciPy v1.6.0开始，cKDTree与KDTree在功能上是相同的，并具有相匹配的性能（请参见注释）。 - Gabriel

Answer 2

尝试使用KDTree.query_ball_point代替。它接受单个点或点数组，并生成距离输入点一定距离内的点。

您可以使用此函数执行批量查询。每次给它100000个点，然后将结果写入文件。类似这样：

BATCH_SIZE = 100000
for i in xrange(0, len(pts), BATCH_SIZE):
    neighbours = tree.query_ball_point(pts[i:i+BATCH_SIZE], X)
    # write neighbours to a file...

- nneonneo

---

除非我理解错了，我认为这正是我列出的潜在解决方案#2，不是吗？据我所知，这种方法的问题在于它需要很长时间。 - Dlinet

你所建议的是循环遍历每一个点。而我所建议的是使用“批处理”模式，这样你就可以花费更少的时间进行迭代。 - nneonneo

啊，有趣，我会研究一下的。我以前从未使用过“批处理”。你建议学习更多关于批处理的特定资源吗？ - Dlinet