时间戳

首页 > 解决方案 > 按距离过滤坐标的最有效方法

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

问题描述

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

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)

输出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])]

除了我需要这样做的数组要大得多。（但只有大约 5% 会低于距离限制，如示例所示）我怎样才能最有效地做到这一点？

标签： pythonarraysnumpy

解决方案

做到这一点的最佳方法之一是使用出色的scipy.spatial.KDTree. 它在114 毫秒内a完成了1000 点和b200K 点的全部工作。这包括：构建 kd-tree、查询它以及过滤半径内的距离。b

（重要：使用scipy >= 1.6.0，其中的 Python 实现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()withdistance_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)

推荐阅读