how to get groups from dataframe when all 1s in flag column and keep group with size greater than 2-CodePudding

i want to extract groups where flag is consective 1 , and keep group if size is greater than 2

here is my df :
df2=pd.DataFrame({'A':[1,20,40,45,56,1,20,40,45,56],'flag':[3,2,4,1,1,3,3,1,1,1]})
print(df2)
    A  flag
0   1     3
1  20     2
2  40     4
3  45     1
4  56     1
5   1     3
6  20     3
7  40     1
8  45     1
9  56     1

output
7  40     1
8  45     1
9  56     1

CodePudding user response：

There are many approaches. You and use inbuilt functions of pandas like groupby and where or simply use the below.

print(df2[(df2['flag']==1) & (df2['A']>2)])

CodePudding user response：

One way using pandas.DataFrame.groupby:

s = df2["flag"].eq(1)
m = s.diff(1).ne(0).cumsum()
new_df = df2[s.groupby(m).transform(lambda x: x.sum()>2)]

Output:

    A  flag
7  40     1
8  45     1
9  56     1