如何实现动态列数DataFrame的列组合求和并追加新列
动态列DataFrame两两组合求和通用实现
核心思路是用标准库的组合生成工具自动枚举所有不重复的两列组合,全程不需要硬编码列数,可直接适配任意列数的数据集。
依赖导入
import pandas as pd from itertools import combinations
通用函数代码
def add_pairwise_sum_cols(input_df: pd.DataFrame) -> pd.DataFrame: # 基于原数据创建副本,避免修改原始DataFrame result_df = input_df.copy() # 自动生成所有长度为2的列组合,列数变化时无需修改逻辑 for col_a, col_b in combinations(input_df.columns, 2): # 新列名直接关联两个源列,方便后续溯源 sum_col_name = f"{col_a}+{col_b}" result_df[sum_col_name] = input_df[col_a] + input_df[col_b] return result_df
方案特性
- 无硬编码适配任意列数:输入3列时自动生成3个求和列,输入N列时自动生成
N*(N-1)/2个两两求和列,完全匹配动态数据集的需求 - 无副作用:所有计算在DataFrame副本上完成,不会改动原始输入数据
- 易扩展:如果需要计算3列、4列甚至更多列的组合和,只需要把
combinations(input_df.columns, 2)里的第二个参数改成对应组合长度即可 - 可溯源:新列名自带求和源列信息,不需要额外维护列对应关系
效果验证
对应题目中3列DataFrame的场景,运行以下测试代码:
# 构造3列测试数据 test_df = pd.DataFrame({ "a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9] }) output_df = add_pairwise_sum_cols(test_df) print(output_df)
得到的输出完全符合预期:
a b c a+b a+c b+c 0 1 4 7 5 8 11 1 2 5 8 7 10 13 2 3 6 9 9 12 15
注意事项
如果数据集包含非数值类型列,建议提前通过
input_df = input_df.select_dtypes(include='number')筛选出数值列再传入函数,避免加法运算报错;如果需要自定义求和列的命名规则,直接修改sum_col_name的生成逻辑即可。
内容的提问来源于stack exchange,提问作者Jacob Mendelowitz
相关产品推荐
相关产品推荐

