解决Pandas中DataFrame非零列关联多类别合并的行结果缺失问题

现有两个DataFrame（规格分别为500×100和1300×2），需要在第一个DataFrame中新增一列，将每行非零值对应的列名关联到第二个DataFrame的category类别，每行可能对应多个类别。

示例数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "apple": [0, 0, 1, 0], 
    "strawberries": [0, 1, 1, 0], 
    "cucumber": [1, 1, 0, 0], 
    "hawthorn": [0, 1, 0, 1]
})

df2 = pd.DataFrame({
    "storage": ["apple", "strawberries", "cucumber", "hawthorn"],
    "category": ["fruits", "berries", "vegetables", "berries"]
})

已尝试的错误方法 #

以下两种基于字典映射的方法仅能得到约1/10的行结果：

df2_dict = dict(zip(df2['storage'], df2['category']))
# 方法一：未按行分组，无法对应每行的多个类别
df['categories'] = pd.Series(df.columns[np.where(df!=0)[1]]).map(df2_dict)

# 方法二：将列名拼成字符串后用map，但字典键是单个列名，无法匹配多列名的字符串
df['categories'] = df.apply(lambda s: ', '.join(s.index[s.eq(1)]), axis = 1).map(df2_dict)

期望输出 #

df = pd.DataFrame({
    "apple": [0, 0, 1, 0], 
    "strawberries": [0, 1, 1, 0], 
    "cucumber": [1, 1, 0, 0], 
    "hawthorn": [0, 1, 0, 1],
    "categories": [
        "vegetables", 
        "berries, vegetables, berries",
        "fruits, berries", 
        "berries"
    ]
})

注：第一个DataFrame的列均为哑变量（仅含0/1值）。

方法一：apply结合列表推导（直观易读） #

先构建存储名到类别的字典，再对每行遍历非零列名，映射后拼接成字符串：

df2_dict = dict(zip(df2['storage'], df2['category']))

df['categories'] = df.apply(
    lambda row: ', '.join([df2_dict[col] for col in row.index if row[col] == 1]),
    axis=1
)

方法二：向量化操作（高效适配大数据集） #

对于500×100的数据集，向量化操作比逐行apply更高效：

# 将df转置后关联category，再转回原结构
category_mapped = df.T.join(df2.set_index('storage')['category']).T

# 筛选每行非零对应的category并拼接
df['categories'] = category_mapped.apply(
    lambda row: ', '.join(row[row != 0].dropna().tolist()),
    axis=1
)

方法三：melt+groupby（结构化重塑实现） #

通过重塑数据结构完成关联与分组拼接：

# 将df转为长格式，保留行索引
df_melted = df.reset_index().melt(id_vars='index', var_name='storage', value_name='flag')
# 筛选非零行并关联category
df_melted = df_melted[df_melted['flag'] == 1].merge(df2, on='storage')
# 按行索引分组，拼接category
category_groups = df_melted.groupby('index')['category'].apply(', '.join)
# 合并回原DataFrame
df = df.join(category_groups)

以上三种方法均可正确生成每行对应的多类别拼接字符串，其中方法二、三更适合处理较大规模的数据集，避免逐行遍历的性能开销。

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