python - 根据条件迭代 pandas 数据框中行子集的最快方法

Question

我有一个包含一列 ID 和一列值的数据框 - 每个 ID 与两行或更多行相关联。当给定 ID 内的任何值子集满足我的条件(在本例中，值的子集净为零)时，我想使用新 ID 来标记这些行，该新 ID 由原始 ID 以及附加数字组成，以唯一地标识这些行。识别子集。

到目前为止我尝试过的代码:

import pandas as pd

d = {'ID': ['2016/01 100','2016/01 100','2016/01 100','2016/01 100','2016/01 200',\
'2016/01 200','2016/01 200','2016/01 200'], 'Value': [-343.68, 343.68, -55.2, 55.2,\
-158.77, 158.77, 123.3, -123.3]}

df = pd.DataFrame(data=d)

df['Cumulative_Sum'] = round(df.Value.cumsum(),2)

print(df)

current_ID = df.loc[0, 'ID']
sum_counter = 1
counter = 0
for row in df.index.values:
    if (df.loc[row, 'Cumulative_Sum'] == 0):
        df.loc[counter:row, 'New_ID'] = str(df.loc[row, 'ID']) + "_" + str(sum_counter)
        counter = row + 1
        sum_counter = sum_counter + 1
        if (counter < len(df.index.values)):    
            if (df.loc[counter, 'ID'] != df.loc[row, 'ID']):
                sum_counter = 1

print (df)

这会产生所需的结果，但在数十万行上运行时速度很慢。

            ID   Value  Cumulative_Sum         New_ID
0  2016/01 100 -343.68         -343.68  2016/01 100_1
1  2016/01 100  343.68            0.00  2016/01 100_1
2  2016/01 100  -55.20          -55.20  2016/01 100_2
3  2016/01 100   55.20            0.00  2016/01 100_2
4  2016/01 200 -158.77         -158.77  2016/01 200_1
5  2016/01 200  158.77            0.00  2016/01 200_1
6  2016/01 200  123.30          123.30  2016/01 200_2
7  2016/01 200 -123.30            0.00  2016/01 200_2

有没有一种更快的方法可以在不循环的情况下完成此操作，同时仍将新 ID 保持为我需要的格式？

Answer 1

我使用嵌套调用来分组，其中使用 cumsum 来查找以零结尾的连续组。最后，使用 ngroup 获取标签。

然后我使用 pd.Series.str.cat 将结果附加到旧 ID。

df.assign(
    New_ID=
    df.ID.str.cat(
        df.groupby('ID').apply(
            lambda d: d.groupby(
                d.Cumulative_Sum.eq(0).iloc[::-1].cumsum(),
                sort=False).ngroup()
        ).add(1).astype(str),
        sep='_'
    )
)

            ID   Value  Cumulative_Sum         New_ID
0  2016/01 100 -343.68         -343.68  2016/01 100_1
1  2016/01 100  343.68            0.00  2016/01 100_1
2  2016/01 100  -55.20          -55.20  2016/01 100_2
3  2016/01 100   55.20            0.00  2016/01 100_2
4  2016/01 200 -158.77         -158.77  2016/01 200_1
5  2016/01 200  158.77            0.00  2016/01 200_1
6  2016/01 200  123.30          123.30  2016/01 200_2
7  2016/01 200 -123.30            0.00  2016/01 200_2