Splitting dataframe into chunks based on negative or np.nan value in specific column

user_10459366 0
Viewed 80

Let's say I have the following dataframe:

import pandas as pd
import numpy as np

df = pd.DataFrame({'dif_seq': [np.nan, 1, 1, 1, 1, -23, 1, 1, 1, -4, 1, 1], 'data': range(12)})

df
Out[75]: 
    dif_seq  data
0       NaN     0
1       1.0     1
2       1.0     2
3       1.0     3
4       1.0     4
5     -23.0     5
6       1.0     6
7       1.0     7
8       1.0     8
9      -4.0     9
10      1.0    10
11      1.0    11

I would like to split df into a list of dataframes based on the values in df['dif_seq'] as follows (all negative or np.nan values signify the start of a new df):

    dif_seq  data
0       NaN     0
1       1.0     1
2       1.0     2
3       1.0     3
4       1.0     4

    dif_seq  data
5     -23.0     5
6       1.0     6
7       1.0     7
8       1.0     8

    dif_seq  data
9      -4.0     9
10      1.0    10
11      1.0    11

What is the best way to go about this? I have an analagous problem with a very large dataset. So although this is a small example, what would be the fastest route to go?

pandas dataframe python split
3 Answers

I would like to split df into a list of dataframes

You can try with a conditional cumulative sum and np.split:

c = df['dif_seq'].lt(0)|df['dif_seq'].isna()
#c= ~df.dif_seq.ge(0) : courtesy @MustafaAydın
s = c.cumsum()
l = np.split(df,np.where(np.diff(s)>0)[0]+1)
#or for a dictionary: dict(iter(df.groupby(s)))

>>l

[   dif_seq  data
 0      NaN     0
 1      1.0     1
 2      1.0     2
 3      1.0     3
 4      1.0     4,
    dif_seq  data
 5    -23.0     5
 6      1.0     6
 7      1.0     7
 8      1.0     8,
     dif_seq  data
 9      -4.0     9
 10      1.0    10
 11      1.0    11]
user_98406370

Create the sub-groupby key with diff and cumsum

s = df['dif_seq'].diff()
s = (s.notnull()& s.ne(0)).cumsum()
s
0     0
1     0
2     0
3     0
4     0
5     1
6     2
7     2
8     2
9     3
10    4
11    4
Name: dif_seq, dtype: int32
d = {x : y for x , y in df.groupby(s)}
user_79645270

Create a series than assigned a chunk and then use it as a mask.

df = pd.DataFrame({'dif_seq': [np.nan, 1, 1, 1, 1, -23, 1, 1, 1, -4, 1, 1], 'data': range(12)})

s = (df["dif_seq"].isna() | df["dif_seq"].lt(0)).cumsum()

split = {f"df{i}":df.loc[s.eq(i)] for i in s.unique()}

split

output

{'df1':    dif_seq  data
 0      NaN     0
 1      1.0     1
 2      1.0     2
 3      1.0     3
 4      1.0     4,
 'df2':    dif_seq  data
 5    -23.0     5
 6      1.0     6
 7      1.0     7
 8      1.0     8,
 'df3':     dif_seq  data
 9      -4.0     9
 10      1.0    10
 11      1.0    11}
user_94414040
Related