Pandas DataFrame与字典互转：还原含NaN的初始DataFrame

没问题，我来帮你拆解这两个需求——先解决从字典反向还原原始带NaN的DataFrame的问题，再优化正向转字典的代码～

首先明确反向转换的核心规则：每个字典键对应原始DataFrame里的一组数据，组内行数是该键下col2和col3列表长度的最大值；col1只有第一行是键值，其余为NaN；col2和col3按各自的列表填充，长度不够的位置在末尾补NaN。

基于这个逻辑，我们可以写一个直观的转换函数：

import pandas as pd
import numpy as np

def dict_to_raw_df(data):
    rows = []
    for val, cols in data.items():
        # 提取两个列的列表，默认空列表防止键缺失
        col2_vals = cols.get("col2", [])
        col3_vals = cols.get("col3", [])
        # 计算当前组需要的总行数
        total_rows = max(len(col2_vals), len(col3_vals))
        
        # 补全两个列到指定长度，末尾补NaN
        col2_full = col2_vals + [np.nan] * (total_rows - len(col2_vals))
        col3_full = col3_vals + [np.nan] * (total_rows - len(col3_vals))
        # 生成col1列：首行是键值，其余为NaN
        col1_full = [val] + [np.nan] * (total_rows - 1)
        
        # 拼接当前组的所有行数据
        for c1, c2, c3 in zip(col1_full, col2_full, col3_full):
            rows.append({"col1": c1, "col2": c2, "col3": c3})
    
    # 转换为DataFrame并返回
    return pd.DataFrame(rows)

# 用你示例中的data测试
test_data = {
    "A1": {"col2": ["B1", "B2", "B3"], "col3": ["C1"]},
    "A2": {"col2": ["B4", "B5"], "col3": ["C2", "C3"]},
    "A3": {"col2": ["B6"], "col3": ["C4", "C5"]}
}

raw_df = dict_to_raw_df(test_data)
print(raw_df)

运行后就能得到和你原始结构完全一致的带NaN的DataFrame。

你原来的字典推导式其实已经可用，但我们可以利用pandas内置的agg方法让代码更简洁高效，尤其是处理大数据集时，agg是pandas内部优化过的方法，比手动循环更节省资源：

方法1：简洁的字典推导式（小数据集友好） #

# 填充col1后的df
df_filled = df[["col1"]].fillna(method="ffill")

# 优化后的字典推导式，逻辑更清晰
data_optimized = {
    val: {
        "col2": group["col2"].dropna().tolist(),
        "col3": group["col3"].dropna().tolist()
    } 
    for val, group in df_filled.groupby("col1")
}

方法2：用agg批量处理（大数据集更高效） #

# 利用agg一次性完成分组聚合，代码更紧凑
data_agg = df_filled.groupby("col1").agg(
    col2=("col2", lambda x: x.dropna().tolist()),
    col3=("col3", lambda x: x.dropna().tolist())
).to_dict("index")

这个方法直接通过agg对每个分组的col2和col3应用dropna转列表的操作，最后转成以col1值为键的字典，和你原来的结果结构完全一致，但代码可读性和执行效率都更高。

你可以测试一下：先用优化后的代码把填充后的df转成字典，再用反向函数转回去，得到的结果和原始带NaN的df完全一致（pandas中np.nan和pd.NA表示空值的差异不影响逻辑）。

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