Question

我正在将代码从MATLAB转换为python，以加快简单操作。我编写了一个包含嵌套循环和条件语句的函数;循环的目的是在与数组y比较时返回数组x中最近元素的索引列表。我按照1e5项目的顺序进行比较，大约需要30秒才能运行。任何有助于加快这一过程的帮助将不胜感激！我使用numba-pro自动即时编译器已经取得了部分成功：

@autojit()
def find_nearest(x,y,idx):
    idx_old = 0
    rng1 = range(y.shape[0])
    rng2 = range(x.shape[0])
    for i in rng1:
        prev = abs(x[idx_old]-y[i])
        for j in rng2:
            if abs(x[j]-y[i]) < prev:
                prev = abs(x[j]-y[i])
                idx_old = j
        idx[i] = idx_old
    return idx

很抱歉这样的菜鸟，我是蟒蛇新手！

Answer 1

你的Numba代码没什么问题，除了算法效率不高。更好的方法是对x数组进行排序并进行二分搜索，与this answer和this answer非常相似：

def find_nearest(x, y):
    indices = np.argsort(x)

    loc = np.searchsorted(x[indices], y)
    right = indices.take(loc, mode='clip')
    left = indices.take(loc-1, mode='clip')

    return np.where(abs(y-x[left]) < abs(y-x[right]), left, right)

在我的电脑上，这比x和y分别有10 ⁶和10 ⁵元素的KDTree方法快约80倍。大约三分之二的时间花费在argsort阵列上，所以我不认为你可以在这里获得Numba。

Answer 2

我找到了解决问题的临时解决方案。通过实现scipy.spatial的kdtree，我能够将运行时间从32秒减少到不到10秒。这仍然比MATLAB knnsearch算法慢四倍;并了解如何使用条件语句加速循环仍然很重要。但目前这一修订后的实施速度更快：

from scipy import spatial
from numpy import matrix

tree = spatial.KDTree(matrix(x).T)
(_, idxx) = tree.query(matrix(y).T)

阵列x和y采用平面1d格式;树需要查询以列向量形式。

非常感谢任何改善原始实施的运行时间的建议！

使用条件语句加速Python嵌套循环

2 个答案: