比pandas iterrows更好的选择

Question

我在pandas中有一个表格，这个表格包含了时间和产品的价格。

为了分析，我想要增加两列，分别显示下一个价格变化超过100美元的时间，无论是上涨还是下跌。

比如说，如果我在09:19这个单元格，接下来价格上涨超过100美元的时间是14:02，而下跌超过100美元的时间是11:39，那么14:02和11:39就应该分别出现在09:19这一行的对应列中。

同样的，如果是09:56这个单元格，接下来价格上涨超过100美元的时间还是14:02，而下跌超过100美元的时间是12:18，这两个值就会出现在09:56这一行。

Table
Time        Price    Up_Time   Down_Time
09:19:00    3252.25     
09:24:00    3259.9      
09:56:00    3199.4      
10:17:00    3222.5      
10:43:00    3191.25     
11:39:00    3143        
12:18:00    2991.7      
13:20:00    3196.35     
13:26:00    3176.1      
13:34:00    3198.85     
13:37:00    3260.75     
14:00:00    3160.85     
14:02:00    3450        
14:19:00    3060.5      
14:30:00    2968.7      
14:31:00    2895.8      
14:52:00    2880.7      
14:53:00    2901.55     
14:55:00    2885.55     
14:57:00    2839.05     
14:58:00    2871.5      
15:00:00    2718.95     

我使用的代码可以实现这个功能，但处理一个数据集需要15到20分钟。

for i, row in df.iterrows():
    time_up = np.nan
    time_down = np.nan

    for j in range(i+1, len(df)):
        diff = df.iloc[j]['Price'] - row['Price']
        if diff > 100:
            time_up = df.iloc[j]['Time']
        elif diff < -100:
            time_down = df.iloc[j]['Time']

        if not pd.isna(time_up) or not pd.isna(time_down):
            break

    df.at[i, 'Up_Time'] = time_up
    df.at[i, 'Down_Time'] = time_down

有没有更高效的方法来做到这一点呢？

Answer 1

你需要把每一行的 Price 值和它后面的所有行进行比较，所以需要进行一些循环操作。你可以使用 apply 和一个函数，结合 numpy 来找到第一个满足变化要求（大于100或小于-100）的值：

def updown(row, df):
    rownum = row.name
    up = (row['Price'] < df.loc[rownum:, 'Price'] - 100).argmax()
    down = (row['Price'] > df.loc[rownum:, 'Price'] + 100).argmax()
    return (
        df.loc[up + rownum, 'Time'] if up > 0 else pd.NaT,
        df.loc[down + rownum, 'Time'] if down > 0 else pd.NaT
    )

df[['Up_Time', 'Down_Time']] = df.apply(updown, axis=1, result_type='expand', df=df)

输出结果：

        Time    Price   Up_Time Down_Time
0   09:19:00  3252.25  14:02:00  11:39:00
1   09:24:00  3259.90  14:02:00  11:39:00
2   09:56:00  3199.40  14:02:00  12:18:00
3   10:17:00  3222.50  14:02:00  12:18:00
4   10:43:00  3191.25  14:02:00  12:18:00
5   11:39:00  3143.00  13:37:00  12:18:00
6   12:18:00  2991.70  13:20:00  14:52:00
7   13:20:00  3196.35  14:02:00  14:19:00
8   13:26:00  3176.10  14:02:00  14:19:00
9   13:34:00  3198.85  14:02:00  14:19:00
10  13:37:00  3260.75  14:02:00  14:19:00
11  14:00:00  3160.85  14:02:00  14:19:00
12  14:02:00  3450.00       NaT  14:19:00
13  14:19:00  3060.50       NaT  14:31:00
14  14:30:00  2968.70       NaT  14:57:00
15  14:31:00  2895.80       NaT  15:00:00
16  14:52:00  2880.70       NaT  15:00:00
17  14:53:00  2901.55       NaT  15:00:00
18  14:55:00  2885.55       NaT  15:00:00
19  14:57:00  2839.05       NaT  15:00:00
20  14:58:00  2871.50       NaT  15:00:00
21  15:00:00  2718.95       NaT       NaT