如何按关联列的多类别归属拆分pandas DataFrame

实现代码 #

方案1：先筛选符合条件的商品列表，再拆分数据 #

这个方法逻辑简单易读，性能表现更好：

import pandas as pd
import numpy as np

# 你的示例构造代码（修正注释语法）
df = pd.DataFrame({'Product':['Prod1','Prod2','Prod3','Prod2','Prod5','Prod3']*4, 
                  'Inv_Type': ['X', 'Y']*12,
                 'Quant': np.random.randint(2,20, size=24)})
df.sort_values('Product', inplace=True, ignore_index=True) # 用于可视化查看

# 第一步：筛选出Inv_Type同时包含X、Y的Product
qualified_prod = df.groupby('Product')['Inv_Type'].nunique()
qualified_prod = qualified_prod[qualified_prod == 2].index.tolist()

# 第二步：拆分数据集
df1 = df[df['Product'].isin(qualified_prod)]
df2 = df[~df['Product'].isin(qualified_prod)]

如果你的Inv_Type可能存在X、Y之外的取值，可以把筛选逻辑换成更严谨的版本：

qualified_prod = df.groupby('Product')['Inv_Type'].apply(lambda x: {'X','Y'}.issubset(set(x)))
qualified_prod = qualified_prod[qualified_prod].index.tolist()

方案2：用transform直接生成判断标识 #

如果不想单独存储符合条件的商品列表，可以直接在原表生成判断列后拆分：

df['is_df1'] = df.groupby('Product')['Inv_Type'].transform(lambda x: {'X','Y'}.issubset(set(x)))
df1 = df[df['is_df1']].drop('is_df1', axis=1).reset_index(drop=True)
df2 = df[~df['is_df1']].drop('is_df1', axis=1).reset_index(drop=True)

两种方案运行后得到的df1、df2和你给出的期望输出完全一致。

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