python - 加快 Numpy/Python 中的最小搜索速度

Question

我有两个 float 数组，想要找到在一定范围内匹配的数据点。这是我到目前为止得到的:

import numpy as np

for vx in range(len(arr1)):
    match = (np.abs(arr2-arr1[vx])).argmin()
    if abs(arr1[vx]-arr2[match])<0.375:
        point = arr2[match]

问题是 arr1 包含 150000 个元素，arr2 包含大约 110000 个元素。这需要花费大量的时间。您有加快速度的建议吗？

Answer 1

除了不进行矢量化之外，您当前的搜索是 (n * m)，其中 n 是 arr2 的大小，m 是 arr1 的大小。在此类搜索中，它有助于对 arr1 或 arr2 进行排序，以便您可以使用二分搜索。排序最终是最慢的步骤，但如果 m 很大，它仍然会更快，因为 n*log(n) 排序比 (n*m) 更快。

以下是如何使用排序数组以矢量化方式进行搜索:

def find_closest(A, target):
    #A must be sorted
    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return A[idx]

arr2.sort()
closest = find_closest(arr2, arr1)
closest = np.where(abs(closest - arr1) < .375, closest, np.nan)