Homogeneous splitting of pandas df

Question

import pandas as pd
import numpy as np

df = pd.DataFrame({'ser_no': [1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
                'co_nm': ['aa', 'aa', 'aa', 'bb', 'bb', 'bb', 'bb', 'cc', 'cc', 'cc', 'aaa', 'aaa', 'aaa', 'bba', 'bba', 'bba', 'bba', 'cca', 'cca', 'cca'],
                'lat': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})

df_splits = np.array_split(df, 4)

This is how I split dataframe. I want the split to be such that no unique 'co_nm' value are spread out i.e. say all the row for which 'co_nm'== 'aa' should be in one split. same goes for all values of 'co_nm'.

Is this possible? existing method can't do it seems.

I think the OP closed it as a dupe, because Community is the only one who closed it. — , Dec 17 '21 at 18:20

score 1 · Answer 1 · 2021-12-17T18:01:33.833

Try using DataFrame.groupby, which will group rows by the unique values of the specified column(s):

splits = [sub_df for _, sub_df in df.groupby('co_nm')]

Output:

>>> splits
[   ser_no co_nm  lat
 0       1    aa    1
 1       2    aa    2
 2       3    aa    3,
     ser_no co_nm  lat
 10       1   aaa    1
 11       2   aaa    2
 12       3   aaa    3,
    ser_no co_nm  lat
 3       4    bb    4
 4       5    bb    5
 5       6    bb    6
 6       7    bb    7,
     ser_no co_nm  lat
 13       4   bba    4
 14       5   bba    5
 15       6   bba    6
 16       7   bba    7,
    ser_no co_nm  lat
 7       8    cc    8
 8       9    cc    9
 9       0    cc   10,
     ser_no co_nm  lat
 17       8   cca    8
 18       9   cca    9
 19       0   cca   10

Use `[df for _, df in df.groupby('co_nm')]` – Dec 17 '21 at 18:00 — , Dec 17 '21 at 18:00

score 1 · Answer 2 · answered Dec 17 '21 at 17:59

You can use groupby on co_nm:

As dict:

d = {k: v for k,v in df.groupby('co_nm')}

As list:

l = [v for _,v in df.groupby('co_nm')]

Output:

{'aa':    ser_no co_nm  lat
0       1    aa    1
1       2    aa    2
2       3    aa    3, 

 'aaa':     ser_no co_nm  lat
10       1   aaa    1
11       2   aaa    2
12       3   aaa    3, 

 'bb':    ser_no co_nm  lat
3       4    bb    4
4       5    bb    5
5       6    bb    6
6       7    bb    7, 

 'bba':     ser_no co_nm  lat
13       4   bba    4
14       5   bba    5
15       6   bba    6
16       7   bba    7, 

 'cc':    ser_no co_nm  lat
7       8    cc    8
8       9    cc    9
9       0    cc   10, 

 'cca':     ser_no co_nm  lat
17       8   cca    8
18       9   cca    9
19       0   cca   10}

score 1 · Answer 3 · answered Dec 17 '21 at 18:02

Split on rows where co_nm changes:

df_splits = np.array_split(df, df.index[df["co_nm"].ne(df["co_nm"].shift())][1:])

>>> df_splits
[   ser_no co_nm  lat
 0       1    aa    1
 1       2    aa    2
 2       3    aa    3,
    ser_no co_nm  lat
 3       4    bb    4
 4       5    bb    5
 5       6    bb    6
 6       7    bb    7,
    ser_no co_nm  lat
 7       8    cc    8
 8       9    cc    9
 9       0    cc   10,
     ser_no co_nm  lat
 10       1   aaa    1
 11       2   aaa    2
 12       3   aaa    3,
     ser_no co_nm  lat
 13       4   bba    4
 14       5   bba    5
 15       6   bba    6
 16       7   bba    7,
     ser_no co_nm  lat
 17       8   cca    8
 18       9   cca    9
 19       0   cca   10]

Homogeneous splitting of pandas df

3 Answers3