为另一个数组中的所有浮点数查找数组中最近的浮点数

Question

根据在另一个数组中找到的最接近的浮点数“过滤”一个数组时，我遇到了性能问题。

这是MWE的问题：

import numpy as np

def random_data(N):
    # Generate some random data.
    return np.random.uniform(0., 10., N).tolist()

# Data lists.
N1 = 1500
list1 = [random_data(N1), random_data(N1), random_data(N1)]
list2 = random_data(1000)

# Define list1's range.
min_1, max_1 = min(list1[2]), max(list1[2])

# This list will contain the "filtered" list1.
list4 = [[], [], []]

# Go through each element in list2.
for elem2 in list2:

    # If it is located within the list1 range.
    if min_1 <= elem2 <= max_1:

        # Find the closest float in sub-list list1[2] to this float
        # in list2.
        indx, elem1 = min(enumerate(list1[2]), key=lambda x:abs(x[1]-elem2))

        # Store the values in list1 that are associated with the closest float
        # found above.
        list4[0].append(list1[0][indx])
        list4[1].append(list1[1][indx])
        list4[2].append(elem1)

（注意list2 包含的元素少于list1[2]，这是我比较的子列表）

此块按预期工作，但效率极低。我确信答案在于正确应用 broadcasting 和 numpy 数组，但我仍然没有设法掌握前者足以将其应用于我的问题。

由于我在增强此代码的性能之后，任何解决方案都可以做到（即：我不受必然使用广播的答案的约束）

---

添加

作为参考，在我前段时间Fast weighted euclidean distance between points in arrays 提出的类似问题中，用户ali_m 使用广播实现了惊人的性能提升。

问题并不完全相同（欧几里得距离而不是绝对值，而且该问题中的距离也必须加权），但在我看来，这个问题比那个问题更简单。

应用于该问题的广播解决方案 ali_m 不能应用于此问题吗？

---

添加 2 个

user2357112 给出的答案以及 Eelco Hoogendoorn 的更正对我最初定义的代码非常有效。我刚刚意识到我过度简化了它，在我的实际代码中，列表list1[2] 和list2 不一定在同一范围内定义。这将是一个更准确的表示（这应该替换上面MWE 中的第一行）：

def random_data(N, xi, xf):
    # Generate some random data.
    return np.random.uniform(xi, xf, N).tolist()

# Data lists.
N1 = 1500
list1 = [random_data(N1, 13., 20.), random_data(N1, -1., 4.), random_data(N1, 2., 7.)]
list2 = random_data(1000, 0., 10.)

现在list1[2] 的范围不等于list2 的范围，因此给出的答案无法拒绝i 的list2[i] > max(list1[2]) 或list2[i] < min(list1[2]) 的那些点。

可以修改答案以考虑这种可能性吗？很抱歉把原来的代码改成这样，真的被我弄丢了。

Answer 1

Kd-tree 在这里真的是大材小用了，你需要做的就是对数组进行排序并使用二分查找在排序后的数组中找到最接近的值。不久前我写了一篇answer，介绍了如何使用searchsorted 来查找数组中目标的壁橱值。您可以在这里使用相同的想法：

import numpy as np

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 idx

def random_data(shape):
    # Generate some random data.
    return np.random.uniform(0., 10., shape)

def main(data, target):
    order = data[2, :].argsort()
    key = data[2, order]
    target = target[(target >= key[0]) & (target <= key[-1])]
    closest = find_closest(key, target)
    return data[:, order[closest]]

N1 = 1500
array1 = random_data((3, N1))
array2 = random_data(1000)
array2[[10, 20]] = [-1., 100]

array4 = main(array1, array2)

Answer 2

如果你有 SciPy，scipy.spatial.cKDTree 可以完成这项工作：

import numpy
import scipy.spatial

array1 = numpy.array(list1)
array2 = numpy.array(list2)

# A tree optimized for nearest-neighbor lookup
tree = scipy.spatial.cKDTree(array1[2, ..., numpy.newaxis])

# The distances from the elements of array2 to their nearest neighbors in
# array1, and the indices of those neighbors.
distances, indices = tree.query(array2[..., numpy.newaxis])

array4 = array1[:, indices]

k-d 树是为多维数据设计的，因此这可能不是最快的解决方案，但与您所拥有的相比，它应该非常快。 k-d 树需要二维点数组形式的输入，其中data[i] 是表示第i 个点的一维数组，因此使用newaxis 的切片表达式将数据放入该格式。如果您需要它更快，您可以使用numpy.sort 和numpy.searchsorted 做一些事情。

如果您需要拒绝来自list2 且超出list1[2] 给出的值范围的数据，可以通过预处理步骤来完成：

lowbound = array1[2].min()
highbound = array1[2].max()

querypoints = array2[(array2 >= lowbound) & (array2 <= highbound)]
distances, indices = tree.query(querypoints[..., numpy.newaxis])