如何使用 python pandas 中的 str.contains 拆分具有多個類別的數據幀？

Question

我有一個像這樣的 dataframe，

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

該表在col1中有 5 個類別。

我知道如何使用 str.contains() 對每個類別進行 select。

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")]

如何使用一些 pandas function 一次生成 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 版本 < 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')