如何将硬编码的Python团队组合生成函数重构为支持任意k个团队的通用函数

核心用递归回溯替代固定层数的硬编码循环，支持任意长度的队伍人数配置，重构后代码如下：

from itertools import combinations

def generate_team_combinations(people, teams):
    # 对每一套队伍人数配置单独生成组合
    for team_config in teams:
        # 内部递归生成器：参数为剩余待分配的人、当前要处理的队伍人数配置切片、已选好的队伍列表
        def backtrack(remaining_people, remaining_slots, selected_teams):
            # 所有队伍都分配完成，返回当前结果
            if not remaining_slots:
                yield selected_teams
                return
            # 取当前要分配的队伍需要的人数
            current_size = remaining_slots[0]
            # 生成所有符合人数的组合
            for comb in combinations(remaining_people, current_size):
                # 剔除已选的人，得到下一轮待分配的人员列表
                next_remaining = [p for p in remaining_people if p not in comb]
                # 递归处理下一支队伍
                yield from backtrack(next_remaining, remaining_slots[1:], selected_teams + [list(comb)])
        
        # 启动递归，初始状态是所有人都待分配，处理完整的队伍配置，已选列表为空
        yield from backtrack(people, team_config, [])

逻辑说明 #

• 外层遍历teams中所有的队伍人数配置方案，支持同时处理多套配置
• 内部递归函数backtrack逐层处理每一支队伍的人员选择：
  • 终止条件：没有剩余待分配的队伍名额时，返回当前已生成的完整队伍组合
  • 每轮从剩余人员中选出当前队伍需要的所有可能组合，更新剩余人员列表后进入下一轮递归
• 全程保持生成器特性，和原代码逻辑一致，不会一次性加载所有组合到内存中

使用示例 #

# 测试：6个人分成3支2人队伍
people = [1,2,3,4,5,6]
teams = [[2,2,2]]
for combo in generate_team_combinations(people, teams):
    print(combo)
# 输出示例：[[1,2], [3,4], [5,6]], [[1,2], [3,5], [4,6]] ... 所有合法组合

可选优化 #

如果人员编号都是唯一值，可以把剩余人员判断逻辑换成集合操作，大幅提升运行效率：

# 把next_remaining的生成逻辑替换为如下代码
next_remaining = list(set(remaining_people) - set(comb))
# 注意：如果需要保持人员的选择顺序稳定性，可以保留原列表推导逻辑

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