如何为Pandas统计DataFrame添加非空唯一州名称列表行？

原始数据 #

year    state   var1    var2    
0   2018    CA       NaN     2    
1   2018    TX       1       NaN    
2   2018    FL       NaN     NaN  
3   2018    AL       1       2    
4   2018    AL       NaN     1   
6   2019    CA       NaN     NaN  
7   2019    TX       1       1    
8   2019    FL       NaN     NaN  
9   2019    AL       2       1    
10  2019    AK       2       NaN 

当前输出结果 #

2018     2019
var1
      Number of unique states with at least 1 non-null:   2        3
      Number of respondents with non-null var:            2        3
      Average:                                            1        1
var2
      Number of unique states with at least 1 non-null:   2        2   
      Number of respondents with non-null var:            3        2
      Average:                                            1.5      1

期望输出结果 #

2018         2019
var1
      Number of unique states with at least 1 non-null:   2            3
      Unique states with at least 1 non-null:             [TX, AL]     [TX, AL, AK]
      Number of respondents with non-null var:            2            3
      Average:                                            1            1
var2
      Number of unique states with at least 1 non-null:   2            2   
      Unique states with at least 1 non-null:             [AL, CA]     [TX, AL]
      Number of respondents with non-null var:            3            2
      Average:                                            1.5          1

修改后的代码 #

import pandas as pd
import numpy as np

# 收集每个变量、年份下的非空唯一州名称列表
state_stats = []
for var in ['var1', 'var2']:
    # 筛选当前变量非空的行，去重后按年份分组提取州名列表
    non_null_state_groups = df[df[var].notna()].drop_duplicates(['year', 'state']).groupby('year')['state'].agg(list)
    # 转换为带多级索引的DataFrame，匹配原有结果的索引结构
    temp_df = non_null_state_groups.to_frame().T
    temp_df.index = pd.MultiIndex.from_tuples([(var, 'Unique states with at least 1 non-null:')])
    state_stats.append(temp_df)

# 合并州名称统计结果
state_df = pd.concat(state_stats)

# 保留原有统计逻辑，修正原代码的大小写错误（YEAR改为year）
c = df.groupby(['year', 'state']).count()
res = c.groupby('year').agg([np.count_nonzero, sum]).T
res.index = res.index.set_levels(['Number of unique states with at least 1 non-null:', 
                                  'Number of respondents with non-null var:'], level=1)

# 计算平均值并整理索引
z = res.swaplevel().T
avg_df = pd.concat([z['Number of respondents with non-null var:'] / z['Number of unique states with at least 1 non-null:']], 
                   keys=['Average:'], axis=1).T.swaplevel().sort_index()

# 合并所有统计结果并排序
final_result = pd.concat([res, state_df, avg_df]).sort_index()
print(final_result)

关键说明 #

1. 州名称统计逻辑：循环遍历目标变量，筛选非空行后去重，按年份分组收集州名列表，转换为与原有结果匹配的多级索引结构。
2. 修正原代码问题：原代码中groupby('YEAR')属于大小写错误，改为groupby('year')避免运行报错。
3. 结果合并：将原有统计数据、州名称列表、平均值统计按多级索引合并后排序，得到符合期望的输出格式。

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