高效求解可重复抽取k个元组的所有唯一逐元素组合和

问题定义 #

从元组列表中可重复抽取k个元组，计算所有唯一的逐元素和（元组求和规则为对应元素相加，例如 (1,2,3)+(4,5,6)=(5,7,9)）。

示例（k=2，n=3） #

输入：

input = [(1,0,0), (2,1,1), (3,3,2)]

所有唯一逐元素和：

{(2, 0, 0), (3, 1, 1), (4, 2, 2), (4, 3, 2), (5, 4, 3), (6, 6, 4)}

实际场景参数 #

• 元组元素范围：0-50（部分维度约束更严，如仅0-2）
• k最大值：4
• 元组长度n最大值：5
• 元组数量最多：1000个

现有算法分析 #

当前采用集合迭代累加的方式，通过每轮迭代生成新的和集合并自动去重，相比直接枚举所有组合（k=4时组合数可达数十亿）更高效，但仍有优化空间：

# 示例元组列表生成
lst = []
for x in range(0,50,2):
    for y in range(0, 20, 1):
        for z in range(0, 3, 1):
            lst.append((x,y,z))

# 通用实现函数
def unique_combination_sums(lst, k):
    n = len(lst[0])
    sums = {tuple(0 for _ in range(n))}  # 初始化为全0元组
    for _ in range(k):
        sums = {tuple(s[i]+x[i] for i in range(n)) for s in sums for x in lst}
    return sums

unique_combination_sums(lst, 4)

该方法的核心问题是每轮迭代需执行集合的笛卡尔积操作，时间复杂度为O(S*M)（S为当前唯一和数量，M为元组数量），当M和k较大时，性能仍会受限。

位集优化方案 #

利用**位集（大整数）**的快速运算特性，将元组的多维和转换为一维整数的位运算，可大幅提升计算效率。

核心思路 #

1. 多维转一维映射：将每个元组映射为唯一整数，利用各维度的最大可能值计算基数，确保不同元组/和对应不同整数。
2. 位集表示集合：用整数的二进制位标记存在的元组/和，例如整数1<<v表示值为v的元组存在。
3. 位运算计算和：位集的乘法等价于所有可能的和的卷积，迭代k次乘法即可得到k个元组相加的所有可能和的位集。

优化后代码实现 #

def unique_combination_sums_bitset(lst, k):
    if not lst:
        return set()
    n = len(lst[0])
    # 元组去重，减少后续计算量
    lst = list(set(lst))
    
    # 计算每个维度k次相加后的最大值
    max_per_dim = [max(t[i] for t in lst) * k for i in range(n)]
    # 计算各维度的基数，用于元组与整数的双向转换
    bases = [1] * n
    for i in range(1, n):
        bases[i] = bases[i-1] * (max_per_dim[i-1] + 1)
    
    # 元组转唯一整数
    def tuple_to_int(t):
        res = 0
        for idx, val in enumerate(t):
            res += val * bases[idx]
        return res
    
    # 生成所有元组对应的整数列表
    tuple_ints = [tuple_to_int(t) for t in lst]
    # 构建元组位集：每个元组对应二进制的一个置位
    tuple_bits = 0
    for v in tuple_ints:
        tuple_bits |= 1 << v
    
    # 初始位集：仅包含全0元组对应的整数0
    current_bits = 1
    # 迭代k次，通过位集乘法得到所有k元组相加的和的位集
    for _ in range(k):
        current_bits *= tuple_bits
    
    # 整数转元组
    def int_to_tuple(num):
        t = []
        remaining = num
        for idx in reversed(range(n)):
            base = bases[idx]
            val = remaining // base
            t.append(val)
            remaining = remaining % base
        return tuple(reversed(t))
    
    # 提取所有置位对应的元组，得到唯一和集合
    result = set()
    num = current_bits
    idx = 0
    while num:
        if num & 1:
            result.add(int_to_tuple(idx))
        num >>= 1
        idx += 1
    return result

# 测试示例
input_lst = [(1,0,0), (2,1,1), (3,3,2)]
print(unique_combination_sums_bitset(input_lst, 2))

性能优势 #

• 运算效率：Python大整数乘法经过底层优化，远快于集合的笛卡尔积操作，尤其当唯一和数量较多时，位集的紧凑存储和批量运算优势明显。
• 内存占用：位集用二进制位标记存在性，相比存储大量元组的集合，内存占用大幅降低。
• 扩展性：对于k=4的场景，位集乘法的迭代成本远低于集合迭代，可轻松处理1000个元组的规模。

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