如何用pd.melt结合正则选择列实现DataFrame宽表转长及后缀处理

问题需求 #

有一个约100列的宽格式DataFrame，需将所有符合XYZ_rating命名格式的列转换为长表格式，同时保留其余所有列不变，并且要移除category列中的_rating后缀。

用户错误代码 #

pd.melt(df, id_vars=str.contains('[^rating]'), value_vars=re.contains(`rating`), var_name='category', value_name='value')

示例DataFrame #

import pandas as pd

df = pd.DataFrame({
    'id': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'first_rating': [1, 2, 3, 1, 2, 3, 1, 2, 3],
    'second_rating': [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
    'third_rating': [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9],
})

期望输出 #

两种正确实现方法 #

方法1：使用pd.melt分步处理 #

先筛选出需要保留的标识列和需要转换的数值列，再转长表后移除后缀：

# 筛选标识列（所有不以_rating结尾的列）
id_cols = df.columns[~df.columns.str.endswith('_rating')]
# 筛选需要转长表的列（所有以_rating结尾的列）
value_cols = df.columns[df.columns.str.endswith('_rating')]

# 执行宽转长
melted_df = pd.melt(df, id_vars=id_cols, value_vars=value_cols, var_name='category', value_name='rating')

# 移除category列的_rating后缀
melted_df['category'] = melted_df['category'].str.replace('_rating', '', regex=False)

方法2：使用pd.wide_to_long一步到位 #

pd.wide_to_long是Pandas专门为这类带固定后缀的宽表设计的工具，代码更简洁：

# 筛选标识列
id_cols = df.columns[~df.columns.str.endswith('_rating')]

# 直接转长表，自动拆分category和rating列
long_df = pd.wide_to_long(df, stubnames=['rating'], i=id_cols, j='category', sep='_', suffix=r'\w+')

# 重置索引，将原索引列转为普通列
long_df = long_df.reset_index()

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