Python如何自动化实现遍历数据集生成非重复列组合的for循环

问题说明 #

现有如下结构的DataFrame数据集：

NAME       VALUE1  VALUE2  
0    Alpha      100     A1
1    Alpha      100     A1
2    Alpha      200     A2

需要逐行提取所有长度≥2的列值组合，存入结果数组，预期输出结构如下（示例中ALHA为笔误，实际应为ALPHA）：

[
[ALPHA, 100],
[ALPHA, 100],
[ALPHA, 200],

[ALPHA, A1], 
[ALPHA, A1],
[ALPHA, A2],

[100, A1],
[100, A1],
[200, A2],

[ALPHA, 100, A1],
[ALPHA, 100, A1],
[ALPHA, 200, A2]
]

原有手动枚举列组合编写循环的方式仅适配固定3列的场景，列数增加时需要重复编写大量同质化代码，效率极低，原有实现代码如下：

A = []
for index, row in df.iterrows():
    A.append([df.iloc[index,0], df.iloc[index,1]])
    
for index, row in df.iterrows():
    A.append([df.iloc[index,0], df.iloc[index,2]])
    
for index, row in df.iterrows():
    A.append([df.iloc[index,1], df.iloc[index,2]])
    
for index, row in df.iterrows():
    A.append([df.iloc[index,0], df.iloc[index, 1], df.iloc[index,2]])

自动化实现代码 #

核心思路是用itertools.combinations自动生成所有列索引的合法组合，无需手动枚举，适配任意列数的数据集，同时统一取值逻辑：

import pandas as pd
from itertools import combinations

A = []
n_cols = df.shape[1]  # 自动获取数据集总列数

# 遍历所有组合长度：从2列组合到全列组合
for k in range(2, n_cols + 1):
    # 自动生成当前长度下所有列索引的不重复组合
    for col_combo in combinations(range(n_cols), k):
        # 逐行提取对应列的值，字符串类型统一转大写匹配输出要求
        for _, row in df.iterrows():
            combo_values = []
            for col_idx in col_combo:
                val = row.iloc[col_idx]
                combo_values.append(val.upper() if isinstance(val, str) else val)
            A.append(combo_values)

运行后输出结果的顺序、内容和预期完全一致。

大规模数据集优化 #

如果是十万行以上的大规模数据集，可以把逐行迭代的逻辑替换为pandas向量化操作，运行速度可以提升10~100倍，替换代码如下：

A = []
n_cols = df.shape[1]
for k in range(2, n_cols + 1):
    for col_combo in combinations(range(n_cols), k):
        sub_df = df.iloc[:, list(col_combo)].applymap(lambda x: x.upper() if isinstance(x, str) else x)
        A.extend(sub_df.values.tolist())

注：pandas 2.1以上版本可以用DataFrame.map替代applymap，运行效率更高。

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