求Pandas列位运算重叠计算的高效优化方案

问题背景 #

现有基于Pandas的位运算重叠百分比计算函数，通过嵌套循环遍历events_x和events_y的所有组合，计算每组事件的重叠占比（sum(x&y)/sum(x|y)）。但在数十万行、数百列的数据集上，循环方式效率极低，需要无遍历的高效实现。

原函数代码：

def get_simple_overlap(dataframe, events_x, events_y):
    df_dict = dict()

    for evt_x, evt_y in product(events_x, events_y):
        overlap = (dataframe[evt_x] & dataframe[evt_y]).tolist()
        total = (dataframe[evt_x] | dataframe[evt_y]).tolist()
        try:
            percentage = sum(overlap) / sum(total)
        except ZeroDivisionError:
            percentage = 0

        if df_dict.get(str(evt_x)) is None:
            df_dict[str(evt_x)] = dict()
        
        df_dict[str(evt_x)][str(evt_y)] = percentage
    
    df = pd.DataFrame(df_dict)

    return df

测试用例：

matrix = pd.DataFrame({
    "evt_x": [0, 1, 0, 1, 1, 1, 0, 1, 0, 1],
    "evt_y": [0, 1, 1, 1, 1, 1, 1, 1, 0, 1],
    # 其他列...
})

event_x = ['evt_x']
event_y = ['evt_y']

overlaps = get_simple_overlap(matrix, event_x, event_y)
# 预期 overlaps['evt_x']['evt_y'] = 6/8 = 0.75

优化方案 #

利用矩阵运算+广播机制替代循环，直接基于整列数据批量计算，大幅提升效率：

核心原理 #

对于0/1类型的列：

• sum(x & y) 等价于两列的点积（x.T @ y），即同时为1的行数
• sum(x | y) 等价于 sum(x) + sum(y) - sum(x & y)，即至少一个为1的行数（避免重复统计重叠部分）
• 重叠占比公式可转化为：dot(x,y) / (sum(x) + sum(y) - dot(x,y))，分母为0时结果设为0

高效实现代码 #

import pandas as pd
import numpy as np

def get_fast_overlap(dataframe, events_x, events_y):
    # 提取对应列的数值矩阵（确保是0/1类型）
    mat_x = dataframe[events_x].values.astype(int)
    mat_y = dataframe[events_y].values.astype(int)
    
    # 计算所有x-y组合的点积（sum(x&y)）
    overlap = mat_x.T @ mat_y
    
    # 计算各列的sum值，广播为矩阵以便批量计算分母
    sum_x = mat_x.sum(axis=0).reshape(-1, 1)  # shape: (len(events_x), 1)
    sum_y = mat_y.sum(axis=0).reshape(1, -1)  # shape: (1, len(events_y))
    total = sum_x + sum_y - overlap
    
    # 计算占比，处理分母为0的情况
    with np.errstate(divide='ignore', invalid='ignore'):
        percentage = np.where(total == 0, 0, overlap / total)
    
    # 转换为指定格式的DataFrame
    return pd.DataFrame(percentage, index=events_y, columns=events_x)

优化效果说明 #

• 时间复杂度从O(MN)（M、N为events_x和events_y的长度）降为O(RM + RN + MN)（R为行数），避免了Python级别的循环，完全利用numpy的向量化运算加速
• 数十万行+数百列的场景下，速度提升可达数十甚至上百倍
• 输出格式与原函数完全一致：events_x为列，events_y为索引

验证 #

用测试用例验证：

matrix = pd.DataFrame({
    "evt_x": [0, 1, 0, 1, 1, 1, 0, 1, 0, 1],
    "evt_y": [0, 1, 1, 1, 1, 1, 1, 1, 0, 1],
})

event_x = ['evt_x']
event_y = ['evt_y']

fast_overlaps = get_fast_overlap(matrix, event_x, event_y)
print(fast_overlaps['evt_x']['evt_y'])  # 输出0.75，与预期一致

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