Pandas Groupby后实现指定子集笛卡尔积的技术需求

先明确你的需求：我们有一个包含col1、col2、col3的DataFrame，需要按col1分组后，把每组里col2为B的去重col3值，和col2为C的去重col3值做笛卡尔积，最后整理成指定格式的结果。

原始DataFrame如下：

你期望得到的结果是：

解决方案代码 #

我用Pandas的分组+自定义处理函数来实现，直接上可运行的代码：

import pandas as pd

# 构造原始数据（如果你的df已经存在，跳过这步即可）
data = [
    ["JHGK", "B", "name1"],
    ["JHGK", "B", "name2"],
    ["JHGK", "C", "name3"],
    ["OPDR", "B", "name4"],
    ["ERTH", "B", "name5"],
    ["ERTH", "C", "name6"]
]
df = pd.DataFrame(data, columns=["col1", "col2", "col3"])

# 定义分组处理函数
def process_single_group(group):
    # 提取组内col2=B的去重col3，重命名为n1
    b_values = group[group["col2"] == "B"]["col3"].drop_duplicates().rename("n1")
    # 提取组内col2=C的去重col3，重命名为n2
    c_values = group[group["col2"] == "C"]["col3"].drop_duplicates().rename("n2")
    
    # 如果其中一类没有数据（比如OPDR组只有B没有C），返回空表过滤掉该组
    if b_values.empty or c_values.empty:
        return pd.DataFrame(columns=["col1", "n1", "n2"])
    
    # 执行笛卡尔积操作
    cross_product = b_values.reset_index(drop=True).merge(
        c_values.reset_index(drop=True), 
        how="cross"
    )
    # 补充当前组的col1标识
    cross_product["col1"] = group["col1"].iloc[0]
    # 调整列顺序到期望格式
    return cross_product[["col1", "n1", "n2"]]

# 按col1分组处理，合并结果并重置索引
final_result = df.groupby("col1").apply(process_single_group).reset_index(drop=True)
print(final_result)

代码细节解释 #

• 自定义的process_single_group函数负责处理单个分组：
  • 分别筛选并去重B、C类的col3值，避免重复值影响笛卡尔积结果
  • 加入空值判断，自动过滤掉只有B或只有C的分组（比如示例中的OPDR组）
  • 使用merge(how="cross")实现笛卡尔积，这是Pandas 1.2.0及以上版本的便捷写法；如果你的版本较低，可以用pd.merge(b_values, c_values, how='outer', left_index=True, right_index=True)替代
• 最后通过groupby.apply把所有分组的结果合并，重置索引后就得到了完全符合需求的结果

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