Splitting pandas DF into equal chunks based on column value

user_5927701 0
Viewed 209

I have a sample df which I am trying to create chunks where a given column value is spread equally (almost) among the chunks.

For example, consider the DF:

df = pd.DataFrame([["apple",1],["mango",20],["mango",20],["mango",31],["apple",17],
                  ["apple",17],["mango",17],["mango",17],["mango",17],["mango",17],["mango",17],["mango",17],["mango",17],
                   ["apple",17],["mango",17],["mango",17],["apple",17]
                  ])
df.columns = ["fruit", "count"]
df

if I want to create 2 chunks of this df where the value for fruit column =apple gets split equally among the 2 chunks.

my current version, first create separate Dfs for mango and apple and then create the splits and concat corresponding splits.

Code:

df_apple = df.loc[df["fruit"]=="apple"]
df_apple_splits = np.array_split(df_apple,2)
df_mango = df.loc[df["fruit"]=="mango"]
df_mango_splits = np.array_split(df_mango,2)

df_split1 = pd.concat([df_apple_splits[0],df_mango_splits[0]])
df_split2 = pd.concat([df_apple_splits[1],df_mango_splits[1]])

But is there a more elegant way to do this?

pandas dataframe numpy python
1 Answers

We can try iterating over a groupby on fruit, array_split into 2 DataFrames, then zip to transpose the list of lists of DataFrames, then concat to create a list of DataFrames (which can be unpacked into two variables):

df_split1, df_split2 = [
    pd.concat(lst) for lst in
    zip(*[np.array_split(v, 2) for _, v in df.groupby('fruit')])
]

df_split1:

   fruit  count
0  apple      1
4  apple     17
5  apple     17
1  mango     20
2  mango     20
3  mango     31
6  mango     17
7  mango     17
8  mango     17

df_split2:

    fruit  count
13  apple     17
16  apple     17
9   mango     17
10  mango     17
11  mango     17
12  mango     17
14  mango     17
15  mango     17

DataFrame and imports:

import numpy as np
import pandas as pd

df = pd.DataFrame({
    'fruit': ['apple', 'mango', 'mango', 'mango', 'apple',
              'apple', 'mango', 'mango', 'mango', 'mango',
              'mango', 'mango', 'mango', 'apple', 'mango',
              'mango', 'apple'],
    'count': [1, 20, 20, 31, 17, 17, 17, 17, 17, 17, 17, 17, 17,
              17, 17, 17, 17]
})
user_154978880
Related