时间戳

首页 > 解决方案 > 如何使用列表（例如数据类）将 MultiIndex DataFrame 转换为对象？

python - 如何使用列表（例如数据类）将 MultiIndex DataFrame 转换为对象？

问题描述

首先，假设您有以下 DataFrame。

import pandas as ps

df = ps.DataFrame([
    [0, 'test0', 0, 'sub0', 'one'],
    [0, 'test0', 1, 'sub1', 'two'],
    [1, 'test1', 0, 'sub0', 'one'],
    [1, 'test1', 1, 'sub1', 'two'],
], columns=['id', 'name', 'sub_id', 'sub_name', 'value'])

df = df.set_index(['id', 'sub_id'])

            name sub_name value
id sub_id                      
0  0       test0     sub0   one
   1       test0     sub1   two
1  0       test1     sub0   one
   1       test1     sub1   two

我想将其转换为如下所示的列表对象（这里我们使用数据类）。

from typing import List
from dataclasses import dataclass

@dataclass
class SubObj:
    id: int
    name: str
    value: str

@dataclass
class MainObj:
    id: int
    name: str
    sub_obj: List[SubObj]

输出应如下所示：

result = [
    MainObj(
        id=0,
        name='test0',
        sub_obj=[
            SubObj(
                id=0,
                name='sub0',
                value='one'
            ),
            SubObj(
                id=1,
                name='sub1',
                value='two'
            )
        ]
    ),
    MainObj(
        id=1,
        name='test1',
        sub_obj=[
            SubObj(
                id=0,
                name='sub0',
                value='one'
            ),
            SubObj(
                id=1,
                name='sub1',
                value='two'
            )
        ]
    ),
]

print(result)

[MainObj(id=0, name='test0', sub_obj=[SubObj(id=0, name='sub0', value='one'), SubObj(id=1, name='sub1', value='two')]), MainObj(id=1, name='test1', sub_obj=[SubObj(id=0, name='sub0', value='one'), SubObj(id=1, name='sub1', value='two')])]

我想实现它，以便它输出一个 MainObj 列表，其中包含尽可能短且易于理解的代码。

你知道怎么做吗？

标签： pythonpandas

解决方案

像这样的小列表理解怎么样？

result = [MainObj(
    row[0][0], 
    row[1]['name'], 
    SubObj(
        row[0][1],
        row[1]['sub_name'],
        row[1]['value']
    )
) for row in df.iterrows()]

退货

[MainObj(id=0, name='test0', sub_obj=SubObj(id=0, name='sub0', value='one')),
 MainObj(id=0, name='test0', sub_obj=SubObj(id=1, name='sub1', value='two')),
 MainObj(id=1, name='test1', sub_obj=SubObj(id=0, name='sub0', value='one')),
 MainObj(id=1, name='test1', sub_obj=SubObj(id=1, name='sub1', value='two'))]

更新

刚刚意识到你想要 sub_obj's 作为列表。我认为这将是一个更好的方法：

results = list()
for _, g in df.groupby(level=0):  # Groupby on first index
    results.append(
        MainObj(
            g.index[0][0],  # Get the first index value
            g['name'].iloc[0],
            [SubObj(row[0][1], row[1]['sub_name'], row[1]['value']) for row in g.iterrows()]))  # List comp iterrating over group rows

[MainObj(id=0, name='test0', sub_obj=[SubObj(id=0, name='sub0', value='one'), SubObj(id=1, name='sub1', value='two')]),
 MainObj(id=1, name='test1', sub_obj=[SubObj(id=0, name='sub0', value='one'), SubObj(id=1, name='sub1', value='two')])]

推荐阅读