Python中如何从多个列表生成符合指定条件的所有唯一组合

实现方案 #

核心思路 #

• 先确定合法的元素选取数量分配：因为要求每个列表至少选1个，首先给每个列表分配1个基础名额，剩余待分配名额为 总长度要求 - 列表总数，本题中为 5-3=2。接下来需要将剩余名额分配给各个列表，要求每个列表最终的选取数量不超过自身的元素总数。
• 对每一种合法的数量分配，分别从每个列表中选取对应数量的元素组合，再将不同列表的组合做笛卡尔积拼接，得到符合要求的最终组合。

Python 代码实现 #

from itertools import combinations, product

def generate_valid_combinations(lists: list[list], target_total: int) -> list[tuple]:
    n = len(lists)
    # 每个列表至少选1个，总长度不够直接返回空
    if target_total < n:
        return []
    # 统计每个列表最多可选取的元素数量
    max_counts = [len(lst) for lst in lists]
    remaining = target_total - n
    res = set()

    # 递归生成所有合法的数量分配方案
    def gen_counts(pos, left, current_counts):
        if pos == n:
            if left == 0:
                # 基础名额+分配的额外名额=实际选取数量
                final_counts = [c+1 for c in current_counts]
                # 过滤超出单列表最大长度的分配方案
                if all(final_counts[i] <= max_counts[i] for i in range(n)):
                    yield final_counts
            return
        # 当前位置最多可分配的额外名额
        max_add = min(left, max_counts[pos] - 1)
        for add in range(0, max_add + 1):
            yield from gen_counts(pos + 1, left - add, current_counts + [add])
    
    # 遍历所有合法分配方案生成组合
    for counts in gen_counts(0, remaining, []):
        list_combs = []
        for i in range(n):
            list_combs.append(list(combinations(lists[i], counts[i])))
        # 对不同列表的组合做笛卡尔积拼接
        for items in product(*list_combs):
            # 扁平化处理，sorted避免选取顺序不同导致的重复组合
            combined = tuple(sorted(elem for sublist in items for elem in sublist))
            res.add(combined)
    
    return list(res)

# 测试示例
list1 = ['a','b','c']
list2 = ['d','e']
list3 = ['f','g','h','i']
all_lists = [list1, list2, list3]
result = generate_valid_combinations(all_lists, 5)
# 输出结果
print(f"总共有{len(result)}种符合要求的组合：")
for comb in result:
    print(comb)

补充说明 #

如果需要保留元素从不同列表选取的顺序特征，不需要去重，可以去掉结果拼接时的sorted逻辑，直接将元素按顺序拼接即可。

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