如何基于动态获取的处理单元名称生成Pandas DataFrame？

问题说明 #

需要根据动态获取的处理单元名称生成对应DataFrame，但通过动态赋值的trname筛选时得到空结果，直接硬编码单元名称（如'California'）却能正常返回数据，需实现通用的动态识别方案。

错误根源 #

代码中这一行是核心问题：

trname = print(str(res)[1:-1])

print()函数的返回值是None，导致trname实际为None，用它匹配State列自然找不到对应数据。另外，res = [treatedunit[0]]属于多余操作，无需将单个元素包装成列表后再做字符串处理。

修正方案 #

直接从treatedunit列表中提取唯一的处理单元名称即可：

trname = treatedunit[0]

修正后的完整代码 #

import pandas as pd

# 读取数据，请替换为你的数据文件路径
df = pd.read_csv('california_prop99.csv', 
sep=';', parse_dates=['Year'], index_col='Year')

# 按州和年份排序数据
df = df.sort_values(by=['State', 'Year'])

# 生成单元唯一ID（hash生成的ID非连续自增，若需从1开始的连续ID，可改为df['id'] = df.groupby('State').ngroup() + 1）
df['id'] = df.State.map(hash)

# 获取所有处理单元的名称列表
treatedunit = df[df['treated'] == 1].State.unique().tolist()

# 断言确保只有一个处理单元（若支持多处理单元可删除此断言）
assert len(treatedunit) == 1

# 动态提取处理单元名称
trname = treatedunit[0]

# 生成目标处理单元的DataFrame
df_treat = df[df['State'] == trname]
print(df_treat)

扩展说明 #

• 若后续需支持多个处理单元，可循环遍历treatedunit列表，逐个生成对应DataFrame并存储（如存入字典）
• 如需生成从1开始的连续单元ID，替换ID生成代码为：df['id'] = df.groupby('State').ngroup() + 1

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