你将如何优化这个简短但非常慢的 Python 循环？

Posted 2023-03-12

【中文标题】你将如何优化这个简短但非常慢的 Python 循环？

【英文标题】：How would you optimize this short but very slow Python loop?

【发布时间】：2020-05-01 13:43:21

【问题描述】：

我正在从 R 切换到 Python。不幸的是，我发现虽然某些结构在 R 中几乎可以立即运行，但在 Python 中却需要几秒钟（甚至几分钟）。阅读后，我发现在 pandas 中强烈建议不要使用 for 循环，建议使用其他替代方法，例如矢量化和应用。

在此示例代码中：从从最小值到最大值排序的一列值中，保留在长度为“200”的间隙之后出现的所有值。

import numpy as np
import pandas as pd

#Let's create the sample data. It consists of a column with random sorted values, and an extra True/False column, where we will flag the values we want
series = np.random.uniform(1,1000000,100000)
test = [True]*100000
data = pd.DataFrame('series' : series, 'test':test )
data.sort_values(by=['series'], inplace=True)

#Loop to get rid of the next values that fall within the '200' threshold after the first next valid value
for i in data['series']:
    if data.loc[data['series'] == i,'test'].item() == True:
        data.loc[(data['series'] > i) & (data['series'] <= i+200  ) ,'test' ] = False
#Finally, let's keep the first values after any'200' threshold             
data = data.loc[data['test']==True , 'series']

是否可以将其转换为函数、矢量化、应用或除“for”循环之外的任何其他结构以使其几乎立即运行？

【问题讨论】：

对于这种动态，for 循环似乎是不可避免的。

注意，pandas 中的.apply 不会比循环快，除非您在列上应用矢量化函数。

我认为问题出在data['series'] == i，这是一个 O(n) 操作，data['series'] > i 和 data['series'] <= i+200 也是 O(n)。所以你有一个外部循环（for i in data['series']），它将运行 O(n) 次，并且在循环内你正在执行 O(n) 操作。所以你的算法是O（n ^ 2）。由于输入的大小为 100,000，因此您将执行大约 10^10 次操作，这很棘手。

补充我上面的评论：我猜 R 中的相应算法在排序后只是 O(n)，所以它可能运行得更快。 （O(n) 算法将遍历列的所有索引i，并检查 i 和 i+1 处的值之间的差异）

【参考方案1】：

您可以通过一个简单的单遍算法在系列上使用一个循环来做到这一点；不需要矢量化或类似的东西。在我的机器上需要 33 毫秒，所以不是“瞬时”，而是眨眼，你会错过它。

def first_after_gap(series, gap=200):
    out = []
    last = float('-inf')
    for x in series:
        if x - last >= gap:
            out.append(x)
            last = x
    return out

例子：

>>> import numpy as np
>>> series = sorted(np.random.uniform(1, 1000000, 100000))
>>> from timeit import timeit
>>> timeit(lambda: first_after_gap(series), number=1)
0.03264855599991279

【讨论】：

IIUC，这与 OP 的预期有些不同。这只是将系列中的每个元素与前一个元素进行比较，这是完全可矢量化的。也许将last = x 放入if x - last >= gap。

我确定它是可向量化的，我只是说没有必要优化它，因为简单的实现足够高效（与“秒或分钟”相比）。行为上的具体区别是什么？

我想说的是，它不是 OP想要的。

我明白了，我问行为上有什么不同？不过，我看到您编辑了您的评论以建议进行更改。我已经编辑了答案并更新了时间测量。

对不起，我错过了那部分。你的方法实际上比我的更干净。不同之处在于，人们不知道需要保留/比较哪些值，因此无法向量化此解决方案。

【参考方案2】：

这是我使用while 循环的方法：

head = 0
indexes = []
while head < len(data):
    thresh = data['series'].iloc[head] + 200
    indexes.append(head)
    head += 1
    while head < len(data) and data['series'].iloc[head] < thresh:
        head+=1

# output:
data = data.iloc[indexes]

# double check with your approach
set(data.loc[data['test']].index) == set(data.iloc[indexes].index)
# output: True

上述方法耗时 984 毫秒，而您的方法耗时 56 秒。

【讨论】：

这就完成了。谢谢。我将尝试将此结构用于类似的任务。不幸的是，我的循环需要很多子集。子集太多和附加数据可能会非常缓慢。你会给我什么建议？

【参考方案3】：

searchsorted

您可以找到下一个，而无需遍历所有...。 这应该更快。 正如 cmets 中所指出的，更快取决于数据。

请注意，我使用与 Quang 类似的方法，因为它们是正确的，您必须循环。不同之处在于我使用searchsorted 来查找每个位置的下一个位置，而不是循环遍历每个位置并评估是否应该添加该位置。

a = data.series.to_numpy()
head = 0
indexes = [head]
while head < len(data):
    head = a[head:].searchsorted(a[head] + 200) + head
    if -1 < head < len(data):
        indexes.append(head)

data.iloc[indexes]

              series  test
77193       5.663829  True
36166     210.829727  True
85730     413.206840  True
68686     613.849315  True
88026     819.096379  True
...              ...   ...
13863  999074.688286  True
31992  999276.058929  True
71844  999487.746496  True
84515  999690.104536  True
6029   999891.101087  True

[4761 rows x 2 columns]

【讨论】：

searchsorted 是 O(log(n))，使用 while 循环可能不会比普通线性循环快。真的取决于数据。

你认为什么样的数据会使这个速度变慢？你让我很好奇（-：

也许你必须添加每个位置？

我正想说完全一样的:-)。

我将阈值更改为0.1 而不是（200），它仍然支持我的方法。但我部分相信你的观点。所以现在我要在这上面浪费太多时间了... thx (-: