寻求更优雅的Pandas方法筛选双生产线共生产产品以对比OEE

核心思路：先找出同时在A、B两条生产线生产过的产品，再用这些产品过滤原数据集，完全避免显式循环，代码更简洁高效。

完整优化代码 #

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# 生成测试数据（保留原代码逻辑）
data = {"batch": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109],
        "production line": ["A", "A", "B", "A", "B", "B", "A", "B", "B", "B"],
        "article": ["hammer", "drill", "hammer", "pliers", "pliers", "pliers", "hammer", "hammer", "hammer", "widget"],
        "OEE": np.random.rand(10)}
df = pd.DataFrame(data, columns = ["batch", "production line", "article","OEE"])

# 关键优化步骤：筛选符合条件的产品
# 方法1：适用于仅A、B两条生产线的场景
valid_articles = df.groupby('article')['production line'].nunique()[lambda x: x == 2].index

# 方法2：通用场景（即使新增生产线，也能精准判断是否包含A和B）
# valid_articles = df.groupby('article')['production line'].apply(lambda x: {'A','B'}.issubset(x))
# valid_articles = valid_articles[valid_articles].index

# 过滤原数据集得到目标结果
reduced = df[df['article'].isin(valid_articles)]

# 绘图（保留原代码逻辑）
sns.stripplot(x= reduced["production line"], y=reduced.OEE, data=reduced, size=16, alpha=.2, jitter=True, edgecolor="none")

优化细节说明 #

1. 性能提升：用Pandas内置的分组操作替代O(n²)复杂度的嵌套循环，大数据量下性能优势显著
2. 逻辑清晰：直接聚焦「产品是否在两条线都有生产」的核心条件，代码可读性更强
3. 扩展性好：方法2可轻松适配多条生产线的场景，只需修改判断的生产线集合即可

原代码问题分析 #

原嵌套循环存在以下问题：

• 时间复杂度高，数据量大时运行缓慢
• 手动修改comparable列易出现索引操作错误
• 代码冗余，逻辑分散，难以维护

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