如何按元素和降序生成组合？替代全量排序的高效方案

问题背景 #

itertools.combinations生成的二元组合是字典序（输入非递增排序时），示例输出如下：

>>> from itertools import combinations
>>> for c in combinations([9,8,7,2,2,1], 2):
...     print(c, sum(c))
... 
(9, 8) 17
(9, 7) 16
(9, 2) 11
(9, 2) 11
(9, 1) 10
(8, 7) 15
(8, 2) 10
(8, 2) 10
(8, 1) 9
(7, 2) 9
(7, 2) 9
(7, 1) 8
(2, 2) 4
(2, 1) 3
(2, 1) 3

需要按元素和从大到小生成组合，但全量生成后排序的方式（sorted(combinations(...), key=sum, reverse=True)）在数据量极大时完全不可行。已知条件：

• 固定生成二元组合（r=2）
• 输入可预先处理为非递增序列
• 存在高成本筛选条件，仅需遍历约0.1%的组合即可找到符合要求的最大和组合

解决方案 #

利用输入非递增的特性，直接按和的降序优先级遍历组合，无需生成所有组合再排序。核心逻辑是从最大可能的和开始，逐个检查每个和对应的所有组合，找到符合条件的就停止，大幅减少计算量。

实现思路 #

1. 计算输入数组的最大和（首两个元素之和）与最小和（最后两个元素之和）
2. 从最大和开始，依次向下遍历每个可能的目标和
3. 对每个目标和s，遍历数组元素arr[i]，计算配对所需的补数s - arr[i]，在i+1到数组末尾的范围内查找该补数（保证i<j，避免重复组合）
4. 生成所有满足arr[i]+arr[j]=s的组合，逐个验证筛选条件，一旦找到符合条件的组合立即终止流程

代码实现 #

def generate_pairs_by_sum_desc(arr):
    # 确保输入为非递增序列
    arr = sorted(arr, reverse=True)
    n = len(arr)
    if n < 2:
        return
    
    max_sum = arr[0] + arr[1]
    min_sum = arr[-1] + arr[-2]
    
    # 从最大和到最小和遍历
    for target_sum in range(max_sum, min_sum - 1, -1):
        i = 0
        while i < n:
            current = arr[i]
            complement = target_sum - current
            # 在i+1到末尾找complement
            j = i + 1
            while j < n:
                if arr[j] == complement:
                    yield (current, arr[j])
                    # 跳过连续重复的complement，避免重复生成相同组合
                    while j + 1 < n and arr[j+1] == complement:
                        j += 1
                        yield (current, arr[j])
                # 跳过连续重复的元素，减少循环次数
                while j + 1 < n and arr[j+1] == arr[j]:
                    j += 1
                j += 1
            # 跳过连续重复的current，减少循环次数
            while i + 1 < n and arr[i+1] == current:
                i += 1
            i += 1

# 测试
test_arr = [9,8,7,2,2,1]
for pair in generate_pairs_by_sum_desc(test_arr):
    print(pair, sum(pair))

关键优化 #

• 跳过重复元素：避免对相同元素重复处理，减少无效循环
• 提前终止：一旦找到符合筛选条件的组合，立即停止遍历，无需处理后续更小的和
• 二分查找优化：将内层j的遍历替换为二分查找（利用数组非递增的特性），可将时间复杂度从O(n²)降至O(n log n)，适合超大数据量场景

内容的提问来源于stack exchange，提问作者no step on snek