pandas中如何生成指定长度m的可变列名组合列表

需求描述 #

需要编写入参为Pandas DataFrame和整数m的函数，自动提取DataFrame的所有列名，生成每m个列名为一组的不重复组合，最终返回存储所有组合的嵌套列表。
以测试DataFrame为例：

import pandas as pd
df = pd.DataFrame({"ID" : [1, 2, 3, 4], 
                  "age": [46, 48, 55, 55],
                  "gender": ['female', 'female', 'male', 'male'],
                  "overweight": ['y', 'n', 'y', 'y']},
                  index = [0, 1, 2, 3])     

当m=2时，预期输出为：
[['ID', 'age'],['ID', 'gender'],['ID', 'overweight'],['age', 'gender'], ['age', 'overweight'], ['gender', 'overweight']]
要求支持动态列数的DataFrame、动态合法取值的m。

具体实现 #

直接使用Python标准库itertools.combinations即可，该方法天生用于生成无重复、无顺序的元素组合，不需要手动编写遍历去重逻辑，适配任意列数、任意合法m值的场景。
完整代码如下：

from itertools import combinations
import pandas as pd

def get_column_combinations(df: pd.DataFrame, m: int) -> list:
    cols = df.columns.tolist()
    col_total = len(cols)
    # 增加参数合法性校验，避免非法m值报错
    if not isinstance(m, int) or m < 1 or m > col_total:
        raise ValueError(f"m必须是1到{col_total}之间的整数，当前传入值不合法")
    # 生成组合并转为嵌套列表格式返回
    return [list(item) for item in combinations(cols, m)]

效果验证 #

调用函数传入示例df和m=2：

print(get_column_combinations(df, m=2))

运行输出与预期完全一致：

[['ID', 'age'], ['ID', 'gender'], ['ID', 'overweight'], ['age', 'gender'], ['age', 'overweight'], ['gender', 'overweight']]

• 当m=1时，会返回[['ID'], ['age'], ['gender'], ['overweight']]
• 当m等于DataFrame总列数时，会返回仅包含全列名组合的单元素列表
• 所有组合不会出现重复分组（例如不会同时存在['ID','age']和['age','ID']的重复项）

内容的提问来源于stack exchange，提问作者peter.bucher