使用Python中的itertools生成总和,同时跟踪总和的每个元素?

时间:2017-05-02 22:37:17

标签: python numpy itertools

我有两个数字列表。对于每个列表,我需要计算所有可能的总和,然后比较生成的总和以查找匹配。我需要输出总和以及包含该总和的每个列表中的元素。如果有多个可能的组合来达到给定的总和,我想保留所有组合。

示例输入:

a = [5, 12.5, 20]
b = [4, 13.5, 20]

期望的输出:

x = [17.5, 20, 37.5] #sums that matched, sorted ascending
a1 = [(5, 12.5),(20),(5, 12.5, 20)] #elements of a in solution sums
b1 = [(4, 13.5),(20),(4, 13.5, 20)] #elements of b in solution sums

这是我到目前为止所尝试的,它找到所有可能的组合,然后是所有可能的总和,并使用numpy数组进行比较。这似乎过于复杂,只有当a和b的长度相同时才有效,而它们可能不是。

import numpy as np
import itertools

a = [5, 12.5, 20]
b = [4, 13.5, 20]

acombos = [seq for i in range(len(a), 0, -1) for seq \
in itertools.combinations(a, i)]

asums = map(sum,acombos)

bcombos = [seq for i in range(len(b), 0, -1) for seq \
in itertools.combinations(b,i)]

bsums = map(sum,bcombos)

comboarray = np.array([acombos, asums, bcombos, bsums])

solutionarray = comboarray[:,comboarray[1] == comboarray[3]]
solutionarray = solutionarray[:, np.argsort(solutionarray[1])]

print(solutionarray)

2 个答案:

答案 0 :(得分:0)

对于每个输入列表,创建一个包含键的和的字典。每个总和都包含一个总和的数字集列表。

找到两个输入列表的总和的交集。然后,对于每个相交的和,打印来自不同列表的数字,这些数字加起来为这些公共总和。

import itertools
import collections


def sums_dict(nums):
    d = collections.defaultdict(list)
    for k in range(1, len(nums) + 1):
        for c in itertools.combinations(nums, k):
            d[sum(c)].append(c)
    return d

def print_matching_sums(a, b):
    ad = sums_dict(a)
    bd = sums_dict(b)
    sums = sorted(set(ad).intersection(bd))
    # The tuples are only to make the output a little easier to read
    print('sums = {}'.format(tuple(sums)))
    print('a = {}'.format(tuple(ad[s] for s in sums)))
    print('b = {}'.format(tuple(bd[s] for s in sums)))


a = [5, 12.5, 20]
b = [4, 13.5, 20]
print_matching_sums(a, b)

输出:

sums = (17.5, 20, 37.5)
a = ([(5, 12.5)], [(20,)], [(5, 12.5, 20)])
b = ([(4, 13.5)], [(20,)], [(4, 13.5, 20)])

答案 1 :(得分:0)

import itertools
import numpy as np
a = [5, 12.5, 20]
b = [4, 13.5, 20]
#get all possible combinations for a and b
a_combos = [list(e) for i in range(1,len(a)+1) for e in itertools.combinations(a,i)]
b_combos = [list(e) for i in range(1,len(b)+1) for e in itertools.combinations(b,i)]

#sum all combinations for a and b
a_sums = np.array(map(sum,a_combos))
b_sums = np.array(map(sum,b_combos))

#get indexes where the sum of a a combo equals the sum of a b combo
a_id, b_id = np.where((a_sums[:,None] == b_sums))

#construct the output using the index
x = a_sums[a_id]
a1 = np.array(a_combos)[a_id]
b1 = np.array(b_combos)[b_id]

x
Out[813]: array([ 20. ,  17.5,  37.5])

a1
Out[814]: array([[20], [5, 12.5], [5, 12.5, 20]], dtype=object)

b1
Out[815]: array([[20], [4, 13.5], [4, 13.5, 20]], dtype=object)