从一组数字中查找长度为3的所有组合,使得sum(组合)为0

时间:2014-06-21 15:07:13

标签: python permutation list-comprehension

假设S是一组整数,例如{-4,-2,1,2,5,0}。我想写一个理解,以获得所有三元素元组(i,j,k)的列表,使i,j,k是S和i+j+k == 0的元素。

3 个答案:

答案 0 :(得分:2)

原帖:

就我理解您的问题和评论的答案而言,i,j,k中的任何值都可能是您的集S中的任何值。

>>> [(i,j,k) for i in S for j in S for k in S if i+j+k == 0]
[(0, 0, 0), (0, 2, -2), (0, -2, 2), (1, 1, -2), (1, -2, 1), (2, 0, -2), (2, 2, -4), (2, -4, 2), (2, -2, 0), (-4, 2, 2), (-2, 0, 2), (-2, 1, 1), (-2, 2, 0)]

修改

由于问题已澄清,即无法选择相同的值两次,此答案需要更新。

>>> from itertools import permutations, combinations
>>> S={-4,-2,1,2,5,0}
>>> [x for x in permutations(S,3) if sum(x) == 0]
[(0, 2, -2), (0, -2, 2), (2, 0, -2), (2, -2, 0), (-2, 0, 2), (-2, 2, 0)]
>>> [x for x in combinations(S,3) if sum(x) == 0]
[(0, 2, -2)]

使用permutationscombinations(我还没想出你想要的那些)。

答案 1 :(得分:1)

itertools docs实际上显示了组合的等效python代码

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = range(r)
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

print [ x for x in combinations(s,3) if sum(x)==0]
[(0, 2, -2)]

答案 2 :(得分:1)

假设不能多次使用相同的数字来计算总和:

from itertools import combinations

S = {-4, -2, 1, 2, 5, 0}
zero_sums = [(i, j, k) for i, j, k in combinations(S, 3) if not sum((i, j, k))]
print(zero_sums)  # -> [(0, 2, -2)]

您可以通过打印所有组合及其总和来验证这是否正确:

for i, j, k in combinations(S, 3):
    print('({:2d}, {:2d}, {:2d}) = {:2d}'.format(i, j, k, sum((i, j, k))))

输出:

( 0,  1,  2) =  3
( 0,  1,  5) =  6
( 0,  1, -4) = -3
( 0,  1, -2) = -1
( 0,  2,  5) =  7
( 0,  2, -4) = -2
( 0,  2, -2) =  0
( 0,  5, -4) =  1
( 0,  5, -2) =  3
( 0, -4, -2) = -6
( 1,  2,  5) =  8
( 1,  2, -4) = -1
( 1,  2, -2) =  1
( 1,  5, -4) =  2
( 1,  5, -2) =  4
( 1, -4, -2) = -5
( 2,  5, -4) =  3
( 2,  5, -2) =  5
( 2, -4, -2) = -4
( 5, -4, -2) = -1