如何从含月度x/y列的Pandas DataFrame生成嵌套列表JSON

我有一个包含id、jan_x、jan_y、feb_x、feb_y……dec_x、dec_y列的Pandas DataFrame，希望将其导出为如下结构的JSON：

{
    "id1": [
        [jan_x, feb_x, ..., dec_x],
        [jan_y, feb_y, ..., dec_y]
    ],
    "id2": [
        [jan_x, feb_x, ..., dec_x],
        [jan_y, feb_y, ..., dec_y]
    ]
}

其中顶级键对应DataFrame中id列的值。请问无需自定义解析函数，是否有简洁的实现方式？我尝试直接导出JSON，但无法得到所需的列表结构。

以下是仅包含两个月数据的示例DataFrame（已修正语法错误）：

data = {'id': ['1', '2', '3', '4'],
        'jan_x': [1, 2, 3, 4],
        'jan_y': [5, 6, 7, 8],
        'feb_x': [9, 10, 11, 12],
        'feb_y': [13, 14, 15, 16]}
  
df = pd.DataFrame(data)

期望输出如下：

{
    "1": [
        [1, 9],
        [5, 13]
    ],
    "2": [
        [2, 10],
        [6, 14]
    ],
    "3": [
        [3, 11],
        [7, 15]
    ],
    "4": [
        [4, 12],
        [8, 16]
    ]
}

解决方案 #

可以通过Pandas的列拆分、多级索引和聚合操作实现，无需自定义解析函数，以下是两种简洁实现方式：

方法一：利用多级索引与链式聚合 #

import pandas as pd
import json

# 修正后的示例数据
data = {'id': ['1', '2', '3', '4'],
        'jan_x': [1, 2, 3, 4],
        'jan_y': [5, 6, 7, 8],
        'feb_x': [9, 10, 11, 12],
        'feb_y': [13, 14, 15, 16]}
  
df = pd.DataFrame(data)

# 1. 将id设为索引，拆分列名为(月份, 维度)的多级索引
df = df.set_index('id')
df.columns = df.columns.str.split('_', expand=True)

# 2. 按维度(x/y)聚合月份值为列表，再组合成目标结构
result = df.unstack().groupby(level=1).agg(list).unstack(level=0).apply(list, axis=1).to_dict()

# 3. 输出格式化JSON
print(json.dumps(result, indent=2))

方法二：分步提取维度值（更直观） #

import pandas as pd
import json

df = pd.DataFrame(data)

# 1. 拆分非id列的列名，构建多级索引
non_id_cols = [col.split('_') for col in df.columns if col != 'id']
df.columns = pd.MultiIndex.from_tuples(non_id_cols, names=['month', 'dim'])
df = df.set_index('id')

# 2. 分别提取x、y维度的所有月份值，转为列表
x_lists = df.xs('x', level='dim', axis=1).apply(list, axis=1)
y_lists = df.xs('y', level='dim', axis=1).apply(list, axis=1)

# 3. 组合x、y列表，生成目标字典
result = pd.concat([x_lists, y_lists], axis=1).apply(list, axis=1).to_dict()

# 输出格式化JSON
print(json.dumps(result, indent=2))

关键逻辑说明 #

• 列拆分：将jan_x这类列名拆分为('jan', 'x')的元组，通过多级索引区分月份和维度，为后续分组提取提供便利。
• 维度提取：用xs()方法快速筛选出所有x或y维度的列，再通过apply(list, axis=1)将每行的月份值合并为一个有序列表。
• 结构组合：将x和y的列表按id合并为包含两个子列表的结构，最终转为字典即可完全匹配目标JSON格式。

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