如何在Python中按公司类别互斥划分训练集与测试集

scikit-learn没有直接满足你需求的内置函数，但可以通过以下简单步骤实现同一公司样本全进训练集或全进测试集的拆分：

• 第一步：提取所有唯一公司
  先从数据集中获取不重复的公司列表：

  unique_companies = df['Company'].unique()
  

• 第二步：拆分公司分组
  用train_test_split先对公司列表进行拆分，这样就能保证同一公司的所有样本后续会被分到同一个集合里：

  from sklearn.model_selection import train_test_split
  
  # 可根据需求调整test_size和random_state参数
  train_companies, test_companies = train_test_split(unique_companies, test_size=0.2, random_state=42)
  

• 第三步：筛选得到最终数据集
  根据拆分后的公司列表，从原数据中提取对应的训练集和测试集：

  train_df = df[df['Company'].isin(train_companies)]
  test_df = df[df['Company'].isin(test_companies)]
  

可选：保持输出标签的分层比例 #

如果需要让训练集和测试集的Output标签分布尽量一致，可以先统计每个公司的主导输出标签，再基于这个标签做分层拆分：

# 统计每个公司占比最高的Output标签
company_dominant_label = df.groupby('Company')['Output'].agg(lambda x: x.value_counts().idxmax())
# 按主导标签分层拆分公司
train_companies, test_companies = train_test_split(unique_companies, 
                                                   test_size=0.2,
                                                   stratify=company_dominant_label.loc[unique_companies],
                                                   random_state=42)

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