用groupby查找pandas dataframe中某列的連續出現次數

Question

我有一個 pandas dataframe 看起來像這樣

import pandas as pd
foo = pd.DataFrame({'id': [1,1,1,1,1,2,2,2,2,2],
                    'col_a': [1,1,0,1,0,1,1,1,0,1],
                    'col_b': [0,1,1,0,0,0,1,1,1,0]})

我想創建 2 列：

• col_a_consequent：如果col_a根據id有n 1 1后續出現
• col_c：如果在col_a 處1出現2次后， col_a有1 ， col_b 1

output dataframe 看起來像這樣：對於n=3

foo = pd.DataFrame({'id': [1,1,1,1,1,2,2,2,2,2],
                        'col_a': [1,1,0,1,0,1,1,1,0,1],
                        'col_b': [0,1,1,0,0,0,1,1,1,0],
                        'col_a_consequent': [0,0,0,0,0,1,1,1,0,0],
                        'col_c': [1,1,1,0,0,1,1,1,1,0]})

對於col_a_consequent根據這個問題我可以得到我想要的

n = 3

foo_tmp = foo.query('id == 2')

(foo_tmp.col_a.groupby((foo_tmp.col_a != foo_tmp.col_a.shift()).\
                  cumsum()).transform('size') * foo_tmp.col_a >= n).astype(int)

但我不知道如何對所有id s 使用groupby執行相同的操作

有任何想法嗎？

Answer 1

首先你可以使用：

N = 3
foo['col_a_consequent'] = (foo
                           .groupby(['id', foo['col_a'].ne(foo['col_a'].shift()).cumsum()])
                           ['col_a'].transform('size').ge(N).astype(int)
                          )

output：

   id  col_a  col_b  col_a_consequent2  col_a_consequent3
0   1      1      0                  1                  0
1   1      1      1                  1                  0
2   1      0      1                  0                  0
3   1      1      0                  0                  0
4   1      0      0                  0                  0
5   2      1      0                  1                  1
6   2      1      1                  1                  1
7   2      1      1                  1                  1
8   2      0      1                  0                  0
9   2      1      0                  0                  0