如何向量化Pandas DataFrame中两列组合的groupby操作

需求说明 #

现有一个Pandas DataFrame，需对列['A','B','C','D','E','F','G']的所有两列组合执行groupby操作，计算Feature_1和Feature_2的求和值。当前通过循环实现，但数据量较大时耗时较长，需更高效的向量化优化方法。

示例数据 #

import pandas as pd

d = {'A': [0,1,1,0,0,1,0,0],
     'B': [1,1,0,0,0,1,0,1],
     'C': [0,0,1,1,0,0,1,0],
     'D': [1,1,1,0,0,1,0,0],
     'E': [0,0,0,1,0,0,0,1],
     'F': [0,0,1,1,0,0,1,0],
     'G': [1,0,1,1,0,0,0,0],
     'Feature_1':[1,0,1,1,0,1,1,0],
     'Feature_2':[0,1,1,0,1,0,0,1]}

df = pd.DataFrame(d)
print(df)

输出：

A  B  C  D  E  F  G  Feature_1  Feature_2
0  0  1  0  1  0  0  1          1          0
1  1  1  0  1  0  0  0          0          1
2  1  0  1  1  0  1  1          1          1
3  0  0  1  0  1  1  1          1          0
4  0  0  0  0  0  0  0          0          1
5  1  1  0  1  0  0  0          1          0
6  0  0  1  0  0  1  0          1          0
7  0  1  0  0  1  0  0          0          1

期望输出 #

Col_1 Col_2  Col_1_value  Col_2_value  Feature_1_Sum  Feature_2_Sum
0     A     B            0            0              2              1
1     A     B            0            1              1              1
2     A     B            1            0              1              1
3     A     B            1            1              1              1
0     A     C            0            0              1              2
...  ...   ...          ...          ...            ...            ...
3     E     G            1            1              1              0
0     F     G            0            0              1              3
1     F     G            0            1              1              0
2     F     G            1            0              1              0
3     F     G            1            1              2              1

77 rows × 6 columns

当前循环实现代码 #

from itertools import combinations

cols_list = ['A','B','C','D','E','F','G']

i=0

for comb in combinations(cols_list, 2):
    i=i+1

    comb_list = [comb[0], comb[1]]

    try:
        del df_gp
    except Exception as e:
        pass

    try:
        del df_comb
    except Exception as e:
        pass

    df_gp = (df.groupby(by=comb_list).agg(Feature_1_Sum=('Feature_1','sum'), Feature_2_Sum=('Feature_2','sum')).reset_index())

    df_gp['Col_1'] = comb[0]
    df_gp['Col_2'] = comb[1]

    df_gp['Col_1_value'] = df_gp[comb[0]]
    df_gp['Col_2_value'] = df_gp[comb[1]]

    df_comb = pd.DataFrame()

    cols = ['Col_1','Col_2','Col_1_value','Col_2_value','Feature_1_Sum','Feature_2_Sum']
    df_comb = df_gp[cols]

    if i==1:
        df_agg = df_comb
    else:
        df_agg = pd.concat([df_agg, df_comb])

df_agg

优化方案 #

方法1：利用pd.crosstab高效计算分组求和 + 批量合并结果 #

原循环的主要开销在于多次创建中间DataFrame和反复执行pd.concat，我们可以通过pd.crosstab直接生成每个列对的2x2求和矩阵，再统一转换为目标格式，最后一次性合并所有结果，大幅提升效率。

import pandas as pd
from itertools import combinations

cols_list = ['A','B','C','D','E','F','G']
result_list = []

for col1, col2 in combinations(cols_list, 2):
    # 计算Feature_1的列对求和矩阵
    ct1 = pd.crosstab(df[col1], df[col2], values=df['Feature_1'], aggfunc='sum').fillna(0).astype(int)
    # 计算Feature_2的列对求和矩阵
    ct2 = pd.crosstab(df[col1], df[col2], values=df['Feature_2'], aggfunc='sum').fillna(0).astype(int)
    
    # 合并两个矩阵并转换为长格式
    merged = ct1.stack().reset_index(name='Feature_1_Sum')
    merged['Feature_2_Sum'] = ct2.stack().values
    
    # 重命名列并添加列标识
    merged.columns = ['Col_1_value', 'Col_2_value', 'Feature_1_Sum', 'Feature_2_Sum']
    merged['Col_1'] = col1
    merged['Col_2'] = col2
    
    # 调整列顺序为目标格式
    merged = merged[['Col_1', 'Col_2', 'Col_1_value', 'Col_2_value', 'Feature_1_Sum', 'Feature_2_Sum']]
    result_list.append(merged)

# 一次性合并所有结果
df_agg = pd.concat(result_list, ignore_index=True)
print(df_agg)

方法2：向量化生成所有组合统计量（进阶） #

如果数据量极大，还可以利用numpy的广播机制直接计算所有列对的分组求和，彻底避免Python循环：

import pandas as pd
import numpy as np
from itertools import combinations

cols_list = ['A','B','C','D','E','F','G']
X = df[cols_list].values  # shape (n_rows, 7)
f1 = df['Feature_1'].values.reshape(-1,1)  # shape (n_rows,1)
f2 = df['Feature_2'].values.reshape(-1,1)

# 生成所有列对的索引组合
col_pairs = list(combinations(range(len(cols_list)), 2))
result_data = []

for i,j in col_pairs:
    # 获取当前列对的取值
    x1 = X[:,i]
    x2 = X[:,j]
    # 生成唯一分组键（0-3，对应0&0,0&1,1&0,1&1）
    groups = x1 * 2 + x2
    # 计算每个分组的特征和
    u, idx = np.unique(groups, return_inverse=True)
    f1_sum = np.bincount(idx, weights=f1.flatten(), minlength=4)
    f2_sum = np.bincount(idx, weights=f2.flatten(), minlength=4)
    
    # 转换为目标格式的行数据
    for g in range(4):
        val1 = g // 2
        val2 = g % 2
        result_data.append([cols_list[i], cols_list[j], val1, val2, int(f1_sum[g]), int(f2_sum[g])])

# 转换为DataFrame
df_agg = pd.DataFrame(result_data, columns=['Col_1', 'Col_2', 'Col_1_value', 'Col_2_value', 'Feature_1_Sum', 'Feature_2_Sum'])
print(df_agg)

这个方法利用numpy的底层数组运算，比纯Pandas操作更快，适合超大规模数据集。

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