How to split pandas dataframe into multiple parts based on consecutively occuring values in a column?

user_7841401 0
Viewed 274

I have a dataframe which I am representing in a tabular format below. The original dataframe is a lot bigger in size and therefore I cannot afford to loop on each row.

col1 | col2 | col3
a      x     1
b      y     1
c      z     0
d      k     1
e      l     1

What I want is split it into subsets of dataframes with consecutive number of 1s in the column col3. So ideally I want to above dataframe to return two dataframes df1 and df2

df1

col1 | col2 | col3
a      x     1
b      y     1

df2

col1 | col2 | col3
d      k     1
e      l     1

Is there an approach like groupby to do this? If I use groupby it returns me all the 4 rows in a dataframe with col3==1. I do not want that as I need two dataframes each consisting of consecutively occuring 1s. One method is to obviously loop by the rows and as and when I find a 0, I can return a dataframe but that is not efficient. Any kind of help is appreciated.

pandas-groupby pandas python-3.x group-by python
1 Answers

First compare values by 1, then create consecutive groups by shift and cumulative sum and last in list comprehension with groupby get all groups:

m1 = df['col3'].eq(1)
g = m1.ne(m1.shift()).cumsum()

dfs = [x for i, x in df[m1].groupby(g)]
print (dfs)
[  col1 col2  col3
0    a    x     1
1    b    y     1,   col1 col2  col3
3    d    k     1
4    e    l     1]

print (dfs[0])
  col1 col2  col3
0    a    x     1
1    b    y     1

If also is necessary remove single 1 rows is added Series.duplicated with keep=False:

print (df)
  col1 col2  col3
0    a    x     1
1    b    y     1
2    c    z     0
3    d    k     1
4    e    l     1
5    f    m     0
6    g    n     1 <- removed

m1 = df['col3'].eq(1)
g = m1.ne(m1.shift()).cumsum()

g = g[g.duplicated(keep=False)]
print (g)
0    1
1    1
3    3
4    3
Name: col3, dtype: int32

dfs = [x for i, x in df[m1].groupby(g)]
print (dfs)
[  col1 col2  col3
0    a    x     1
1    b    y     1,   col1 col2  col3
3    d    k     1
4    e    l     1]
user_29010020
Related