关于性能：使用条件语句加速Python嵌套循环


Speed up Python nested loops with conditional statements

我正在将代码从matlab转换为python，以加速简单的操作。我已经编写了一个包含嵌套循环和条件语句的函数；该循环的目的是返回数组x中最近元素的索引列表(与数组y相比)。我按照1e5项的顺序进行比较，运行大约需要30秒。任何帮助加快这一进程将非常感谢！我已经部分成功地使用了numba pro自动实时编译器：

1
2
3
4
5
6
7
8
9
10
11
12
13
@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的新手！

相关讨论

你的numba代码没有什么问题，只是算法没有那么有效。更好的方法是对x数组进行排序并执行二进制搜索，与此答案和此答案非常相似：

1
2
3
4
5
6
7
8
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)

号

在我的个人电脑上，这比使用kdtree方法的x和y分别具有106和105个元素快80倍。大约三分之二的时间是花在阵列上的，所以我认为在这里使用numba不会获得太多好处。

相关讨论

我找到了解决我问题的临时办法。通过实现scipy.spatial的kdree i，可以将运行时间从32秒缩短到10秒以下。这仍然比matlab knnsearch算法慢4倍；了解如何用条件语句加速循环仍然很重要。但目前，这一修订的实施速度更快：

1
2
3
4
5
from scipy import spatial
from numpy import matrix

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

数组x和y采用一维平面格式；树要求查询采用列向量形式。

如有任何改进原实施运行时间的建议，将不胜感激！