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

Question

我正在将代码从MATLAB转换为python，以加快简单的操作。我写了一个包含嵌套循环和条件语句的函数;循环的目的是返回数组x中与数组y比较时最接近的元素的索引列表。我按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

对不起，我是一个新手，我是一个全新的python!

Answer 1

您的Numba代码没有任何问题，只是算法没有达到应有的效率。更好的方法是对x数组进行排序并进行二叉搜索，类似于这个答案和这个答案:

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)

在我的PC上，这比x和y的KDTree方法分别有10⁶和10⁵元素快80倍。大约三分之二的时间都花在了对数组进行argsort处理上，所以我认为在这里使用Numba并不能获得多少好处。

Answer 2

我已经找到了一个解决问题的临时办法。通过实现scipy。我能够将运行时间从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格式;该树要求查询是列向量形式。

任何改善原始实现运行时的建议都将非常感谢!