如何用Python/Pandas按分组值特征拆分原始DataFrame为两个DataFrame

输入数据 #

group   treatment
0   111 lot1
1   111 lot2
2   111 lot3
3   111 lot4
4   111 other
5   111 unknown
6   111 other 
7   111 unknown
8   111 another
9   222 other
10  222 unknown
11  222 unknown
12  222 other
13  333 lot1
14  333 lot2
15  333 lot3
16  333 lot4
17  333 other
18  333 unknown
19  444 other
20  444 unknown
21  444 unknown
22  444 other

需求 #

输入数据包含两类分组：

• 第一类分组（111、333）：treatment列包含lot1-lot4、other、unknown（111还额外有another）
• 第二类分组（222、444）：treatment列仅包含other和unknown

需要将原始DataFrame拆分为两个DataFrame：

• df1：包含所有treatment列存在lot1-lot4的分组（即111、333的所有行）
• df2：仅包含treatment列只有other和unknown的分组（即222、444的所有行）

预期输出 #

df1 #

group   treatment
111 lot1
111 lot2
111 lot3
111 lot4
111 other
111 unknown
111 other 
111 unknown
111 another
333 lot1
333 lot2
333 lot3
333 lot4
333 other
333 unknown

df2 #

group   treatment
222 other
222 unknown
222 unknown
222 other
444 other
444 unknown
444 unknown
444 other

Python/Pandas实现代码 #

import pandas as pd

# 构造原始DataFrame
data = {
    'group': [111,111,111,111,111,111,111,111,111,222,222,222,222,333,333,333,333,333,333,444,444,444,444],
    'treatment': ['lot1','lot2','lot3','lot4','other','unknown','other ','unknown','another','other','unknown','unknown','other','lot1','lot2','lot3','lot4','other','unknown','other','unknown','unknown','other']
}
df = pd.DataFrame(data)

# 按group分组，获取每个分组的唯一treatment值
group_unique_treatments = df.groupby('group')['treatment'].unique()

# 定义lot类别的集合，筛选包含lot1-lot4的分组
lot_set = {'lot1', 'lot2', 'lot3', 'lot4'}
target_groups = group_unique_treatments[
    group_unique_treatments.apply(lambda x: len(set(x) & lot_set) > 0)
].index

# 拆分DataFrame
df1 = df[df['group'].isin(target_groups)].reset_index(drop=True)
df2 = df[~df['group'].isin(target_groups)].reset_index(drop=True)

# 输出结果
print("df1:\n", df1)
print("\ndf2:\n", df2)

代码说明 #

1. 先按group字段分组，提取每个分组下所有唯一的treatment值
2. 通过集合交集判断，筛选出包含lot1-lot4任意一个值的分组
3. 用isin()和反向筛选~isin()拆分原始DataFrame，最后重置索引让结果更整洁

内容的提问来源于stack exchange，提问作者Sir