如何通过列除法为Pandas分组结果新增average列

我创建了一个按年份分组的DataFrame，用于统计至少有1个非空值的唯一州数量，以及非空值的总数。但在新增一列（用非空值总数除以唯一州数量）时遇到了困难。

示例数据 #

nan = float('nan')
df = pd.DataFrame({
    'year': [2018, 2018, 2018, 2018, 2018, 2019, 2019, 2019, 2019, 2019],
    'state': [1, 2, 3, 1, 2, 1, 2, 3, 4, 5],
    'var1': [nan, 1.0, nan, 1.0, nan, nan, 1.0, nan, 2.0, 2.0],
    'var2': [2.0, 1.0, nan, 2.0, 1.0, nan, 1.0, nan, 1.0, nan],
})

当前代码 #

import numpy as np

c = df.groupby(['year', 'state']).count()
res = c.groupby('year').agg([np.count_nonzero, sum])

当前输出 #

var1                    var2
        count_nonzero   sum     count_nonzero   sum
year                
2018    2               2       2               4
2019    3               3       2               2

期望输出 #

var1                                    var2
        count_nonzero   sum     average         count_nonzero   sum    average
year                
2018    2               2       1               2               4      2
2019    3               3       1               2               2      1

方法一：直接计算并插入平均值列 #

利用原结果的多层列结构，分别对var1和var2计算sum / count_nonzero，再将结果作为新层级列添加，最后调整列顺序匹配期望输出：

import numpy as np
import pandas as pd

nan = float('nan')
df = pd.DataFrame({
    'year': [2018, 2018, 2018, 2018, 2018, 2019, 2019, 2019, 2019, 2019],
    'state': [1, 2, 3, 1, 2, 1, 2, 3, 4, 5],
    'var1': [nan, 1.0, nan, 1.0, nan, nan, 1.0, nan, 2.0, 2.0],
    'var2': [2.0, 1.0, nan, 2.0, 1.0, nan, 1.0, nan, 1.0, nan],
})

c = df.groupby(['year', 'state']).count()
res = c.groupby('year').agg([np.count_nonzero, sum])

# 计算并添加var1的平均值列
res[('var1', 'average')] = res[('var1', 'sum')] / res[('var1', 'count_nonzero')]
# 计算并添加var2的平均值列
res[('var2', 'average')] = res[('var2', 'sum')] / res[('var2', 'count_nonzero')]

# 调整列顺序，匹配期望输出的结构
res = res.reindex(columns=['var1', 'var2'], level=0)
res = res.reindex(columns=['count_nonzero', 'sum', 'average'], level=1)

print(res)

方法二：自定义聚合函数 #

通过自定义聚合逻辑，在分组时直接生成包含平均值的统计结果，适合需要扩展更多统计项的场景：

import numpy as np
import pandas as pd

nan = float('nan')
df = pd.DataFrame({
    'year': [2018, 2018, 2018, 2018, 2018, 2019, 2019, 2019, 2019, 2019],
    'state': [1, 2, 3, 1, 2, 1, 2, 3, 4, 5],
    'var1': [nan, 1.0, nan, 1.0, nan, nan, 1.0, nan, 2.0, 2.0],
    'var2': [2.0, 1.0, nan, 2.0, 1.0, nan, 1.0, nan, 1.0, nan],
})

def agg_stats(x):
    count = np.count_nonzero(x)
    total = x.sum()
    # 避免除数为0的情况，返回nan
    avg = total / count if count != 0 else np.nan
    return pd.Series([count, total, avg], index=['count_nonzero', 'sum', 'average'])

# 先按year+state分组统计，再按year聚合生成完整统计项
c = df.groupby(['year', 'state']).count()
res = c.groupby('year').apply(lambda group: group.apply(agg_stats))

print(res)

两种方法都能得到你期望的输出，方法一更直接高效，方法二更灵活。

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