Как разделить кадры данных по нескольким категориям, используя str.contains в пандах Python?

Question

У меня есть такой фрейм данных,

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

В этой таблице 5 категорий в col1.

Я знаю, как выбрать каждую категорию, используя 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")]

Как я могу сгенерировать 5 кадров данных одновременно, используя какую-нибудь функцию панд? Таким образом, мои выходные данные представляют собой 5 фреймов данных с именем df_apple, df_peach, df_melon, df_berry, df_banana.

и сохранили их в 5 различных файлах CSV.

Answer 1

Я бы explode столбец найти уникальный 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

---

Или сбросить в 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")

Затем

pd.read_csv('apple.csv')

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

---

Для версий Pandas <<code>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)

И смотрите этот пост @ MaxU

Answer 2

Рекомендую хранить их в диктовке:

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')