Pandas多列快速向量求和优化：避免DataFrame拷贝提速问询

问题分析 #

原代码的核心瓶颈在于[df[x] for x in combos]步骤：每次从DataFrame中选取列子集都会创建副本（尤其是当列不连续时），大量副本的生成和内存占用会拖慢整体速度。你提到的通过引用方式计算求和完全可行，关键是避免创建不必要的DataFrame子集副本，同时利用数学简化减少计算量。

优化方案 #

• 利用组合互补性减少计算量：你的场景是从100列中选98列求和，这等价于每行的总列和 减去 被排除的2列的和。这个转换能把需要处理的列数从98降到2，计算量和内存占用都会大幅降低。
• 直接操作底层numpy数组避免副本：Pandas DataFrame的底层存储是numpy数组，直接操作数组可以避免创建DataFrame子集副本，利用numpy的高效计算能力提升速度。

优化后的代码 #

import itertools as it
import numpy as np
import pandas as pd

def sum_combo_ranks(df, combos: list[str]) -> pd.DataFrame:
    """Returns the sum of the composite ranks for each combination of columns..

    Returns:
        DataFrame
            Each column contains a single system
            Each row is a date/ticker
            Each cell contains the aggregated rank for that stock on that date.
    """

    # Each DataFrame in lst is the sum of a single combination of selected columns;
    # This step is a bottleneck.
    # 2.9s
    lst = [df[x] for x in combos]

    # For each subset, add each row across.
    # 2.1s
    sums = [x.sum(axis=1)
            for x in lst]

    # Join the list of df's into a single DataFrame.
    # 0.3s
    new_df = pd.DataFrame(sums)

    return new_df

def sum_combo_ranks_optimized(df, combos: list[list[int]]) -> pd.DataFrame:
    # 获取底层numpy数组，避免创建DataFrame副本
    arr = df.to_numpy()
    # 先计算每行的总列和（只算一次）
    total_row_sum = arr.sum(axis=1)
    
    # 转换思路：每个98列组合对应唯一的2列排除组合
    # 生成每个组合对应的排除列（比遍历98列更高效）
    all_cols = set(range(df.shape[1]))
    excluded_combos = []
    for combo in combos:
        excluded = list(all_cols - set(combo))
        excluded_combos.append(excluded)
    
    # 计算每个排除组合的行和，再用总行和减去这个值，得到98列的和
    sums = []
    for excluded in excluded_combos:
        # 直接对numpy数组切片求和，无副本创建
        excluded_sum = arr[:, excluded].sum(axis=1)
        sums.append(total_row_sum - excluded_sum)
    
    # 转置后生成结果DataFrame（与原代码输出结构一致）
    new_df = pd.DataFrame(sums)
    return new_df

# 测试代码
df = pd.DataFrame(np.zeros([6000, 100]))       
combos = list(list(x) for x in it.combinations(range(100), 98))[:1000]  

# 对比原函数和优化函数的速度
import time
start = time.time()
result_old = sum_combo_ranks(df, combos)
print(f"原函数耗时: {time.time() - start:.2f}s")

start = time.time()
result_new = sum_combo_ranks_optimized(df, combos)
print(f"优化后耗时: {time.time() - start:.2f}s")

# 验证结果一致性
print(f"结果是否一致: {result_old.equals(result_new)}")

进一步矢量化优化（可选） #

如果需要处理的组合数量很大，可以用numpy广播一次性计算所有结果，彻底消除Python循环的开销：

def sum_combo_ranks_vectorized(df, combos: list[list[int]]) -> pd.DataFrame:
    arr = df.to_numpy()
    total_row_sum = arr.sum(axis=1)
    
    # 把所有排除组合转换成numpy数组（shape=(N,2)，N为组合数量）
    all_cols = set(range(df.shape[1]))
    excluded_arr = np.array([list(all_cols - set(combo)) for combo in combos])
    
    # 一次性计算所有排除组合的行和：利用广播实现批量运算
    excluded_sums = arr[:, excluded_arr].sum(axis=2)
    # 总行和减去每个排除组合的和，得到所有98列的和
    sums = total_row_sum[:, np.newaxis] - excluded_sums
    
    # 转置后生成结果（与原代码输出结构一致）
    new_df = pd.DataFrame(sums.T)
    return new_df

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