python - 如何在 Pandas 中比较前一行并使用 group by 更改另一列的值？

Question

我有一个像这样的数据框，

>>> import pandas as pd
>>> data = {
    'code': ['123456', '123456', '123456', '123abc', '123abc', '123abc'],
    'price': [112.58, 112.58, 113.25, 120.31, 120.36, 140.12],
}... ... ...
>>> all_df = pd.DataFrame(data)
>>> all_df
     code   price
0  123456  112.58
1  123456  112.58
2  123456  113.25
3  123abc  120.31
4  123abc  120.36
5  123abc  140.12

我想创建一个名为 period 的列，并在价格不相同时增加周期。每笔交易都应根据唯一的代码进行。这是预期的输出，

     code   price  period
0  123456  112.58     1.0
1  123456  112.58     1.0
2  123456  113.25     2.0
3  123abc  112.58     1.0
4  123abc  112.58     1.0
5  123abc  113.25     2.0

我是这样做的，但是当数据太多的时候，时间太长了。

def get_period(df):
    period = df.loc[0, 'period'] = 1
    df = df.reset_index(drop=True)

    for i in range(1, len(df)):
        if df.loc[i, 'price'] != df.loc[i - 1, 'price']:
            df.loc[i, 'period'] = period + 1
            period += 1
        else:
            df.loc[i, 'period'] = period

    return df

all_df.groupby('code').apply(get_period).reset_index(drop=True)

有什么想法吗？提前致谢。

Answer 1

首先通过 Series.shift 和 Series.cumsum 测试连续的 price 值，然后为每个组调用 factorize 的每个值:

df['period'] = (df.assign(new = df['price'].ne(df['price'].shift()).cumsum())
                  .groupby('code')['new']
                  .transform(lambda x: pd.factorize(x)[0]) + 1)
print (df)
     code   price  period
0  123456  112.58       1
1  123456  112.58       1
2  123456  113.25       2
3  123abc  112.58       1
4  123abc  112.58       1
5  123abc  113.25       2