How to split dataframes with multiple categories using str.contains in python pandas?

Question

I have a dataframe like this,

id   col1
1    apple, peach
2    apple, banana
3    melon, peach
4    berry, apple, peach
5    melon, banana

This table has 5 categories in col1 .

I know how to select each category using str.contains().

df_apple = df[df['col1'].str.contains("apple")]
df_peach = df[df['col1'].str.contains("peach")]
df_melon = df[df['col1'].str.contains("melon")]
df_berry = df[df['col1'].str.contains("berry")]
df_banana = df[df['col1'].str.contains("banana")]

How can I generate 5 dataframes in one time using some pandas function? So my outputs are 5 dataframes named df_apple, df_peach, df_melon, df_berry, df_banana .

And saved them into 5 different csv files.

Answer 1

I'd explode the column an find unique id

d = df.set_index('id').col1
e = d.str.split(', ').explode()

r = {k: d.loc[v] for k, v in e.index.groupby(e).items()}

r['apple']

id
1           apple, peach
2          apple, banana
4    berry, apple, peach
Name: col1, dtype: object

---

Or to dump to csv

d = df.set_index('id').col1
e = d.str.split(', ').explode()

for k, v in e.index.groupby(e).items():
    d.loc[v].to_frame().to_csv(f"{k}.csv")

Then

pd.read_csv('apple.csv')

   id                 col1
0   1         apple, peach
1   2        apple, banana
2   4  berry, apple, peach

---

For Pandas versions < 0.25

def explode(s):
    return pd.Series(np.concatenate(s.to_numpy()), s.index.repeat(s.str.len()))

d = df.set_index('id').col1
e = d.str.split(', ').pipe(explode)

And see this post by @MaxU

Answer 2

I recommand you to store them in a dict:

dfdict = {fruit:df[df['col1'].str.contains(fruit)] for fruit in ['apple', 'peach', 'melon', 'berry', 'banana']}

for k,v in dfdict.items():
    v.to_csv('df'+k+'.csv')