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

Question

我正在从 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”循环之外的任何其他结构以使其几乎立即运行？

Answer 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

Answer 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 秒。

Answer 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]