条件下的累计计数重置

Question

我有一个与此类似的数据框：

df = pd.DataFrame({'col1': ['a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'c', 'c', 'c', 'c'],
                 'col2': [1, 1, 1, 1, 2, 2, 1, 1, 2, 1, 1, 2, 2],
                 'col3': [1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0],
                 'desired': [0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1]})

我想应用一个滚动总和col3，当col1或col2更改时重置，或者当之前的值为col3零时。

请注意，计数偏移了 1 个单元格。这意味着新(col1, col2)组合的期望值将始终为零。

下面的代码演示了所需的逻辑。但是，在下面的数据集上需要将近 4 分钟。

des = []
count = 0
for i in range(1, len(df)):
    des.append(count)
    if (df.iloc[i-1].col1 == df.iloc[i].col1) & \
       (df.iloc[i-1].col2 == df.iloc[i].col2) & \
       (df.iloc[i-1].col3 == 1):
    
        count += 1
    else:
        count = 0
    
des.append(0)

df['desired'] = des

要测试的更大数据集：https : //www.dropbox.com/s/hbafcq6hdkh4r9r/test.csv?dl=0

Answer 1

首先使用groupbywith shift，然后连续计数1：

a = df.groupby(['col1','col2'])['col3'].shift().fillna(0).eq(1)
b = a.cumsum()

df['desired'] = b-b.where(~a).ffill().fillna(0).astype(int)

print (df.head(20))
      col1  col2  col3  desired
0   100055     1     1        0
1   100055     1     0        1
2   100055     1     0        0
3   100055     1     0        0
4   100055     1     0        0
5   100055     1     0        0
6   100055     1     0        0
7   100055     1     0        0
8   100055     1     0        0
9   100055     1     0        0
10  100055     1     1        0
11  100055     1     1        1
12  100055     1     0        2
13  100055     1     1        0
14  100055     1     1        1
15  100055     1     0        2
16  100055     1     0        0
17  100055     1     1        0
18  100055     1     0        1
19  100055     1     1        0