查找一個點半徑內的所有點（Python）

Question

我有兩個包含多個點的數據集。 兩個數據集都有坐標為 gps 的點。 該點位於同一區域。

我必須從數據集 A 中找到一個點，從數據集 B 中找到所有點，這些點在數據集 A 的點半徑內。

我可以像個白痴一樣用循環來做，但是對於大數據集來說需要很多時間。

你能給我一些算法提示，這可以幫助我找到性能最好的解決方案嗎？

Answer 1

如果您目前正在使用基本循環來執行計算，那么您很可能無需更改算法就可以大大提高性能，只需通過重構您的計算方式即可。

我運行了一系列基准測試來說明不同的方法如何影響性能。 為簡單起見，我使用簡單的 x/y 坐標而不是 Haversine 公式或類似的東西來設置這些基准，但適用相同的原則。

基准進程中有兩個關鍵思想。 首先，在可能的情況下簡化計算：操作越少，花費的時間就越少。 其次，將循環計算移至 numpy，執行向量運算的開銷要低得多。

不同方法的性能：

方法	時間	相對的
基本循環法	45.039	22.68
檢查 r 平方刪除平方根操作	28.941	14.58
使用列表理解加速列表的構建	26.190	13.19
使用 1d arrays 意味着在 C 向量操作中執行循環	1.986	1.00
使用 2d numpy arrays 可以幫助 memory 位置，但在這種情況下不是	3.452	1.74

這些結果來自以下測試代碼：

import numpy as np
import random
import timeit

random.seed(1)

count_a = 10000000
r = 2.5
r_squared = r * r

a_tuples = [(random.uniform(0, 10), random.uniform(0, 10)) for _ in range(count_a)]
b = (random.uniform(4, 8), random.uniform(4, 8))

# tuple of 1d numpy arrays
a_arrs = (np.array([v[0] for v in a_tuples]), np.array([v[1] for v in a_tuples]))

# 2d numpy arrays
a_2d = np.array(a_tuples)
b_arr = np.array(b)

def method0():
    """Basic loop method"""
    out = []
    for x, y in a_tuples:
        dx = x - b[0]
        dy = y - b[1]
        distance = ((dx * dx) + (dy * dy)) ** .5
        if distance <= r:
            out.append((x, y))
    return out

def method1():
    """Checking against r squared removes square root operation"""
    out = []
    for x, y in a_tuples:
        dx = x - b[0]
        dy = y - b[1]
        if ((dx * dx) + (dy * dy)) <= r_squared:
            out.append((x, y))
    return out

def method2():
    """Using list comprehension speeds construction of list"""
    return [(x, y) for x, y in a_tuples if (((dx := x - b[0]) * dx) + ((dy := y - b[1]) * dy)) <= r_squared]

def method3():
    """Using 1d arrays means loops are performed in C vector operations"""
    x, y = a_arrs
    dx = x - b[0]
    dy = y - b[1]
    return np.nonzero(((dx * dx) + (dy * dy)) <= r_squared)[0]

def method4():
    """Using 2d numpy arrays can help with memory locality, but not in this case"""
    a_b = a_2d - b
    a_b2 = a_b * a_b
    sum_sq = np.sum(a_b2, axis=1)
    return np.nonzero(sum_sq <= r_squared)[0]

m0 = method0()
assert m0 == method1()
assert m0 == method2()
assert m0 == [(a_arrs[0][i], a_arrs[1][i]) for i in method3()]
assert m0 == [(a_2d[i][0], a_2d[i][1]) for i in method4()]

number = 10
t0, t1, t2, t3 = None, None, None, None
t0 = timeit.timeit(lambda: method0(), number=number)
t1 = timeit.timeit(lambda: method1(), number=number)
t2 = timeit.timeit(lambda: method2(), number=number)
t3 = timeit.timeit(lambda: method3(), number=number)
t4 = timeit.timeit(lambda: method4(), number=number)

tmin = min((t0, t1, t2, t3, t4))

print(f'| Method               | Time     | Relative        |')
print(f'|------------------    |------    |---------------  |')
print(f'| {method0.__doc__}    | {t0:.3f} | {t0 / tmin:.2f} |')
print(f'| {method1.__doc__}    | {t1:.3f} | {t1 / tmin:.2f} |')
print(f'| {method2.__doc__}    | {t2:.3f} | {t2 / tmin:.2f} |')
print(f'| {method3.__doc__}    | {t3:.3f} | {t3 / tmin:.2f} |')
print(f'| {method4.__doc__}    | {t4:.3f} | {t4 / tmin:.2f} |')