如何使用pivot_table转换指定结构的Pandas数据集？

需求说明 #

需要将每个id对应的带_drug后缀的列中值为1的列名，去掉后缀后用逗号合并成therapy列，同时去除同一id的重复行。

解决方案 #

先构造示例数据（你可以替换成自己的数据集读取代码）：

import pandas as pd

data = {
    'id': [1,1,1,2,2,3,3,3,4,4,5,5],
    'drug1_drug': [0,0,0,0,0,1,1,1,1,1,1,1],
    'drug2_drug': [1,1,1,1,1,1,1,1,0,0,0,0],
    'drug3_drug': [0,0,0,1,1,0,0,0,1,1,0,0],
    'age': [33,33,33,45,45,66,66,66,28,28,87,87]
}
df = pd.DataFrame(data)

方法1：直观逐行提取（新手友好）

1. 先筛选出所有药物相关的列：

drug_cols = [col for col in df.columns if col.endswith('_drug')]

2. 定义函数，提取每行值为1的药物名称并合并：

def get_therapy(row):
    # 遍历药物列，筛选值为1的列，去掉_drug后缀
    selected_drugs = [col.replace('_drug', '') for col in drug_cols if row[col] == 1]
    return ','.join(selected_drugs)

3. 给原数据添加therapy列：

df['therapy'] = df.apply(get_therapy, axis=1)

4. 去重并保留需要的列：

result = df[['id', 'age', 'therapy']].drop_duplicates().reset_index(drop=True)

方法2：使用pivot_table（符合你的思路）

因为同一id的药物标记和年龄是固定的，我们可以用pivot_table聚合每个id的药物最大值（0或1，最大值就是该id是否使用此药），再生成therapy列：

1. 分组聚合：

pivot_df = df.pivot_table(
    index=['id', 'age'],  # 按id和age分组
    values=drug_cols,     # 聚合药物列
    aggfunc='max'         # 取最大值，确保同一id的药物标记统一
).reset_index()

2. 生成therapy列：

pivot_df['therapy'] = pivot_df.apply(
    lambda row: ','.join([col.replace('_drug', '') for col in drug_cols if row[col]==1]),
    axis=1
)

3. 整理结果：

result = pivot_df[['id', 'age', 'therapy']]

两种方法最终都会得到你需要的输出：

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