如何用Pandas将DataFrame分组列转为指定键名的JSON列表并保留指定列

现有DataFrame结构 #

ID A1 B1 C1 A2 B2 C2 A3 B3 C3
AA  1  3  6           4  0  6 
BB  5  5  4  6  7  9 
CC  5  5  5           

需求说明 #

• 新增列Z，将每行中每3个同前缀列（如A1、B1、C1）转换为键为A、B、C的JSON对象，所有有效对象组成JSON列表
• 处理完成后仅保留ID和Z列，期望输出如下：

ID Z
AA [{"A":1, "B":3,"C":6},{"A":4, "B":0,"C":6}]
BB [{"A":5, "B":5,"C":4},{"A":6, "B":7,"C":9}]
CC [{"A":5, "B":5,"C":5}]

当前尝试代码及问题 #

用户尝试的代码：

df2 = df.groupby(['ID']).apply(lambda x: x[['A1', 'B1', 'C1',
                                            'A2', 'B2', 'C2', 'A3', 'B3', 'C3']].to_dict('records')).to_frame('Z').reset_index()

存在两个问题：

• 无法去除列名中的数字，导致JSON键为A1、B1而非A、B
• 无法将每3列划分为独立的JSON对象，输出结构不符合需求

解决方案 #

可以通过按行遍历处理每组列的方式实现需求，具体代码如下：

import pandas as pd
import numpy as np

# 构造示例DataFrame（匹配用户提供的结构）
data = {
    'ID': ['AA', 'BB', 'CC'],
    'A1': [1, 5, 5],
    'B1': [3, 5, 5],
    'C1': [6, 4, 5],
    'A2': [np.nan, 6, np.nan],
    'B2': [np.nan, 7, np.nan],
    'C2': [np.nan, 9, np.nan],
    'A3': [4, np.nan, np.nan],
    'B3': [0, np.nan, np.nan],
    'C3': [6, np.nan, np.nan]
}
df = pd.DataFrame(data)

# 定义每行的处理函数
def process_row(row):
    json_list = []
    # 遍历3组前缀列（1-3号组）
    for group_num in range(1, 4):
        # 获取当前组的3列
        cols = [f'A{group_num}', f'B{group_num}', f'C{group_num}']
        group_vals = row[cols]
        # 仅保留无缺失值的组（避免空对象）
        if not group_vals.isna().any():
            # 将列名的数字去除，生成目标字典
            json_obj = {col[0]: val for col, val in group_vals.items()}
            json_list.append(json_obj)
    return json_list

# 生成Z列
df['Z'] = df.apply(process_row, axis=1)

# 保留指定列
result_df = df[['ID', 'Z']]
print(result_df)

代码说明 #

1. 构造示例数据：将用户提供的文本表格转换为可操作的Pandas DataFrame，用np.nan表示空值
2. 行处理函数：
   • 遍历3组列（1-3号），每组取对应前缀的3列
   • 跳过存在缺失值的组（避免生成不完整的JSON对象）
   • 将列名的首字符作为JSON键，对应值保留，组成单个JSON对象
   • 收集所有有效JSON对象为列表
3. 生成目标列：通过apply按行执行处理函数，得到Z列
4. 筛选列：仅保留ID和Z列，得到最终结果

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