带有 if 语句的 for 循环运行缓慢

Question

我有一个长数据框，大约有 150k 行和两列（A 和 B）。我正在为“i”行在“B”列中编写代码，如果 A 列中接下来的 6 行 == 1，则 B=1.

到目前为止，我设法编写了一个工作代码（见下文）。但是到运行很慢，大概5min。您对如何提高代码效率有什么建议吗？谢谢！

df[j]['B'] = 0

for i in np.arange(0,len(df[j]['A'])-1, dtype=np.float64):
    try:
        if (df[j]['A'][i+1] == 1) & (df[j]['A'][i+2]==1) & (df[j]['A'][i+3]==1) & (df[j]['A'][i+4]==1) & (df[j]['A'][i+5]==1) & (df[j]['A'][i+6]==1) == True:
            df[j]['B'][i] = 1
    except:
        pass

我已经尝试删除 try/except 块但仍然没有改善，实际上它慢了 1 秒。

for i  in np.arange(0,len(WT_df[j]['ID10'])-7, dtype=np.float64):
    if (WT_df[j]['ID10'][i+1] == 1) & (WT_df[j]['ID10'][i+2]==1) & (WT_df[j]['ID10'][i+3]==1) & (WT_df[j]['ID10'][i+4]==1) & (WT_df[j]['ID10'][i+5]==1) & (WT_df[j]['ID10'][i+6]==1) == True:
        WT_df[j]['ID_on'][i] = 1

此致，路易斯

Answer 1

Pandas 直接索引很慢（至少在循环中）。您可以 向量化循环 以加快速度。这是一个未经测试的例子：

column = WT_df[j]['ID10'].values
size = column.size-7
result = np.ones(size, dtype=bool)

for k in range(1, 7):
    result &= column[k:size+k] == 1

WT_df[j].loc[np.where(result)[0], 'ID_on'] = 1

带有 if 语句的 for 循环运行缓慢

for loop with if statement runs slow

optimization

for-loop

python-3.x