生成器链式调用解决组合分配问题：效率优化与N桶适配

一、当前实现的效率分析 #

你的现有实现能正确生成结果，但存在几个明显的效率瓶颈：

• 重复集合转换：每次循环都把possibilities中的列表重复转为集合（比如set(possibilities['b'])在每个comb_a迭代中都执行一次），这会产生大量不必要的计算开销。
• 无提前剪枝：没有提前判断剩余可用物品数量是否满足当前桶的容量要求，即使可用物品不够，仍会进入combinations循环（而此时combinations会返回空，完全是无效计算）。
• 硬编码桶数量：只能处理固定3个桶，扩展性极差，无法适配N个桶的场景。

针对这些问题，优化后的3桶版本可以这样写：

from itertools import combinations

def solve_optimized(possibilities, sizes):
    # 提前将所有候选列表转为集合，避免重复转换
    poss_sets = {k: set(v) for k, v in possibilities.items()}
    buckets = list(possibilities.keys())
    bucket_a, bucket_b, bucket_c = buckets
    
    for comb_a in combinations(possibilities[bucket_a], sizes[bucket_a]):
        comb_a_set = set(comb_a)
        b_available = poss_sets[bucket_b] - comb_a_set
        # 提前剪枝：可用数量不足则跳过
        if len(b_available) < sizes[bucket_b]:
            continue
        for comb_b in combinations(b_available, sizes[bucket_b]):
            comb_b_set = set(comb_b)
            c_available = poss_sets[bucket_c] - comb_a_set - comb_b_set
            if len(c_available) < sizes[bucket_c]:
                continue
            for comb_c in combinations(c_available, sizes[bucket_c]):
                yield comb_a, comb_b, comb_c

二、处理任意数量桶的通用实现 #

用回溯递归的方式可以轻松实现支持任意N个桶的生成器，核心思路是逐个处理每个桶，传递已选物品集合，递归生成后续桶的可行组合：

from itertools import combinations

def solve_general(possibilities, sizes):
    # 预转换候选列表为集合，减少重复计算
    poss_sets = {k: set(v) for k, v in possibilities.items()}
    # 获取桶的处理顺序（可根据需求调整顺序）
    buckets = list(possibilities.keys())
    
    def backtrack(bucket_idx, selected_items, chosen_combs):
        # 所有桶处理完成，返回当前完整方案
        if bucket_idx == len(buckets):
            yield tuple(chosen_combs)
            return
        
        current_bucket = buckets[bucket_idx]
        required_size = sizes[current_bucket]
        # 计算当前桶的可用物品：候选池减去已被其他桶选走的物品
        available = poss_sets[current_bucket] - selected_items
        
        # 提前剪枝：可用物品不足则直接回溯
        if len(available) < required_size:
            return
        
        # 若需要保持组合物品的顺序与原候选列表一致，可用下面的过滤方式：
        # available_list = [item for item in possibilities[current_bucket] if item in available]
        available_list = list(available)
        
        # 遍历当前桶的所有可行组合，递归处理下一个桶
        for comb in combinations(available_list, required_size):
            yield from backtrack(
                bucket_idx + 1,
                selected_items | set(comb),
                chosen_combs + [comb]
            )
    
    # 从第一个桶开始递归生成方案
    yield from backtrack(0, set(), [])

测试示例 #

possibilities = {'a': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15],
                 'b': [0, 1, 2, 3, 5, 6, 7, 9, 10, 11, 12, 13, 14],
                 'c': [1, 2, 3, 5, 6, 7, 9, 10, 11, 12, 13, 14]}

sizes = {'a': 5, 'b': 5, 'c': 6}

count = 0
for solution in solve_general(possibilities, sizes):
    count +=1
print(f"找到 {count} 个方案")  # 输出：找到 11550 个方案

这个通用版本不仅支持任意数量的桶，还通过预转集合、提前剪枝等方式优化了效率，同时如果需要保持组合中物品的顺序与原候选列表一致，只需替换available_list的生成方式即可。

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