如何修改Python组合查找函数，支持带标识列表输入并返回对应结果

需求背景 #

原有基于itertools的findPairs函数可从数值列表中找出和近似等于指定值K的组合，但现在需要适配带标识的二维列表输入，输出同时返回对应的标识组合与数值组合。

原代码 #

from itertools import takewhile, combinations, permutations, product

def findPairs(lst, K):
    for i in range(1,len(lst),1):
        print([pair for pair in combinations(lst, i) if sum(pair) >= K-0.01 and sum(pair) <= K+0.01])

期望输入与输出 #

• 期望输入（带标识的二维列表）：

lst = [['A1',263.09], ['A2',883.58],['A3', 75.75], ['A4',29.88],['A5',30.36],['A6',30.37],['A7',590.72],['A8', 162.45], ['A9',47.25], ['A10',252.98], ['A11',69.57],['A12', 20.24]]
K = 1282.66

• 期望输出：

[(('A1','A2','A3','A4','A5'), (263.09, 883.58, 75.75, 29.88, 30.36)), (('A1','A2','A3','A4','A6'), (263.09, 883.58, 75.75, 29.88, 30.37))]

修改后的函数方案 #

from itertools import combinations

def findPairs(lst, K):
    result = []
    # 遍历所有可能的组合长度（从1到列表长度）
    for i in range(1, len(lst)+1):
        # 生成当前长度的所有组合
        for combo in combinations(lst, i):
            # 计算组合中数值的和
            total = sum(item[1] for item in combo)
            # 判断和是否在K±0.01范围内
            if K - 0.01 <= total <= K + 0.01:
                # 提取标识元组和数值元组
                ids = tuple(item[0] for item in combo)
                values = tuple(item[1] for item in combo)
                result.append((ids, values))
    return result

# 调用示例
if __name__ == "__main__":
    lst = [['A1',263.09], ['A2',883.58],['A3', 75.75], ['A4',29.88],['A5',30.36],['A6',30.37],['A7',590.72],['A8', 162.45], ['A9',47.25], ['A10',252.98], ['A11',69.57],['A12', 20.24]]
    K = 1282.66
    output = findPairs(lst, K)
    print(output)

关键改动说明 #

1. 组合处理：遍历组合时，每个元素是二维列表中的[标识, 数值]项，不再是单纯数值。
2. 求和逻辑：计算总和时，提取每个组合元素的item[1]（数值部分）进行求和。
3. 结果构造：符合条件的组合，分别提取item[0]（标识）和item[1]（数值）组成元组，存入结果列表。
4. 返回结果：将原函数的直接打印改为返回结果列表，便于后续处理大数据集时的灵活使用。

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