按距离过滤坐标的最有效方法

Question

我有两个二维坐标（x，y）的列表/平面 arrays a和b 。 （ a最多包含 ~1000 个坐标， b最多 ~200,000。）我需要过滤b以仅获取 b 中的那些坐标的列表，这些坐标在b中的任何坐标的a欧几里得距离内，比如 10。 例如：

import numpy as np

np.random.seed(123)

a = np.random.randint(50, size=(5, 2))
b = np.random.randint(50, size=(1000, 2))

def dist(coord_a, coord_b):
    squ_dist_x = (coord_a[0] - coord_b[0])**2
    squ_dist_y = (coord_a[1] - coord_b[1])**2
    return np.sqrt(squ_dist_x + squ_dist_y)

b_filter = []
for coord_b in b:
    for coord_a in a:
        distance = dist(coord_a, coord_b)
        if distance <= 5:
            b_filter.append(coord_b)

Output 用于b_filter ：

[array([21, 30]),
 array([ 3, 30]),
 array([ 4, 30]),
 array([48, 32]),
 array([ 8, 32]),
 array([21, 32]),
 array([16, 30]),
 array([16, 30]),
 array([ 0, 32]),
 array([45, 32]),
 array([32, 30]),
 array([ 3, 31]),
 array([21, 30]),
 array([20, 31]),
 array([ 6, 32]),
 array([24, 31]),
 array([31, 31]),
 array([ 2, 31]),
 array([32, 32]),
 array([22, 32]),
 array([46, 30]),
 array([40, 32]),
 array([20, 31]),
 array([ 4, 32]),
 array([31, 31]),
 array([35, 30]),
 array([37, 32]),
 array([ 1, 32]),
 array([31, 31]),
 array([44, 30]),
 array([24, 32]),
 array([40, 32]),
 array([31, 30]),
 array([ 2, 31]),
 array([47, 32]),
 array([42, 32]),
 array([36, 30]),
 array([ 8, 31]),
 array([19, 30]),
 array([42, 31]),
 array([48, 31]),
 array([ 4, 30]),
 array([47, 32]),
 array([30, 32]),
 array([47, 31]),
 array([29, 31]),
 array([29, 31]),
 array([17, 32]),
 array([24, 32]),
 array([46, 30]),
 array([ 2, 30]),
 array([ 4, 30]),
 array([27, 30]),
 array([ 1, 31]),
 array([30, 32]),
 array([ 5, 32]),
 array([25, 32]),
 array([35, 30]),
 array([23, 31]),
 array([36, 31]),
 array([41, 30]),
 array([17, 30]),
 array([26, 30]),
 array([19, 30]),
 array([22, 31]),
 array([ 2, 31]),
 array([23, 32]),
 array([30, 30]),
 array([ 3, 31]),
 array([40, 30])]

除了 arrays 之外，我需要使用更大的尺寸。 （但只有大约 5% 会低于距离限制，如示例所示）我怎样才能最有效地做到这一点？

Answer 1

最好的方法之一是使用出色的scipy.spatial.KDTree 。 它在114 毫秒内完成了a :1000 点和b :200K 点的全部工作。 这包括：构建 kd-tree，查询它，并过滤b以获取半径内的距离。

（重要：使用scipy >= 1.6.0 ，其中 KDTree 的KDTree实现已被cKDTree替换。请参阅发行说明）。

# please use scipy >= 1.6.0
from scipy.spatial import KDTree

tree = KDTree(a)
radius = 10
dist, idx = tree.query(b, k=1, distance_upper_bound=radius + 1)

b_filtered = b[dist <= radius]

解释

• 我们将tree.query()与distance_upper_bound=radius + 1一起使用，以便包含恰好在给定半径处的点（ .query()的行为是返回严格小于上限的点）。
• 然后我们过滤小于或等于半径的距离。
• 当然，如果您对dist < radius的点感兴趣，则无需添加 1，但这并没有什么坏处，因为您必须在下一步中进行过滤（没有近邻的点将+inf为距离）。

定时

a = np.random.randint(50, size=(1000, 2))
b = np.random.randint(50, size=(200_000, 2))

radius = 10
%timeit b_filtered = b[KDTree(a).query(b, k=1, distance_upper_bound=radius + 1)[0] <= radius]
# 114 ms ± 17 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

改进代码的一种简单方法是在此处避免使用sqrt ：

return np.sqrt(squ_dist_x + squ_dist_y)

只需返回总和，然后与 25 进行比较。

顺便提一句：

这似乎是错误的

    squ_dist_x = (coord_a[0] - coord_b[1])**2
                                       ^
                                      [0] ??

Answer 3

使用广播，您可以创建一个索引，指示哪个坐标比n单位更接近：

# Minimal distance:
n = 5
# Creation of the index:
ind = (((b[None,:]-a[:,None])**2).sum(2)**0.5<=n).sum(0) > 0
# Get the result:
b_filter = b[ind]

其中b[None,:]-a[:,None]返回包含坐标差的(5,1000,2)数组。 然后我们简单地应用欧几里得距离公式，并检查是否至少有一个值a距离b比n更近。

Answer 4

如果出于某种原因，您想在纯 numpy 中执行此操作，这里是一个矢量化（但幼稚）的表达式：

b_filtered = b[np.any(np.square(b[:, None] - a).sum(axis=2) <= radius**2, axis=1)]

但这是O(m * n)对于m ， n点分别在a ， b 。 对于m, n = 1000, 200_000 ，需要超过 6 秒。

我提供的另一个使用scipy.spatial.KDTree的答案要快得多（它减少了很多计算）并且在 114 毫秒内完成。