Python - get all possible sums of elements inside n 1d arrays

user_1067083 0
Viewed 4555

Given an integer n, and an array a, I would like to return an array with all the possible values of sums of a with itself n times.

Example: n = 3, a = [1, 2, 3, 4, 5, 6]

Output: [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]

First element is from 1+1+1, second is 1+1+2 etc.

Is there any elegant way to do that? I've tried loops, but since n isn't known in advance, I don't know how many loops I need to make.

Thanks in advance

python
5 Answers

Generate all possible 3-element combinations, then sum them:

from itertools import combinations_with_replacement

n = 3
li = [1, 2, 3, 4, 5, 6]

print([sum(comb) for comb in combinations_with_replacement(li, n)])

# [3, 4, 5, 6, 7, 8, 5, 6, 7, 8, 9, 7, 8, 9, 10, 9, 10, 11, 11, 12, 13, 6, 7, 8, 9, 10, 8, 9, 10, 11, 10, 11, 12, 12, 13, 14, 9, 10, 11, 12, 11, 12, 13, 13, 14, 15, 12, 13, 14, 14, 15, 16, 15, 16, 17, 18]

Since you seem to be interested in the unique sums, use a set:

print(set(sum(comb) for comb in combinations_with_replacement(li, n)))

# {3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}

Note that there is no guarantee whatsoever that these will be ordered. If you want ordered output be explicit about it:

print(sorted(set(sum(comb) for comb in combinations_with_replacement(li, n))))
user_14538220

An alternative solution could be to use itertools.product. Here you first generate pairs of 3 elements from a and then sum them. To get rid of the duplicates, you use set { } and the summation is done usingn list comprehension. Here I am using *[a]*n to make it more general for any value of n. 

import itertools
n = 3
totals = {sum(item) for item in itertools.product(*[a]*n)}
# {3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}

Another readable way to do it would be to use a variable lists and then pass it to itertools.product

lists = [a]*n
totals = {sum(item) for item in itertools.product(*lists)}

To get all the possible sums including the duplicates, simply use [ ] instead of { }.

user_49323160

This will work for you and give you a set as the output to ensure unique sum values. n and a can be any integer or list, respectively.

import itertools

n = 3
a = [1, 2, 3, 4, 5, 6]
b = [a for _ in range(n)]
sums = set(sum(_b) for _b in itertools.product(*b))
user_104165820

A pure-Python implementation:

def comb_sum(arr, n):
    if n == 1:
        [(yield a) for a in arr]
    else:
        for i, a in enumerate(arr):
            [(yield a + b) for b in comb_sum(arr[i:], n-1)]

my_list = [1, 2, 3, 4, 5, 6]
n = 3
sums = set([c for c in comb_sum(my_list, n)])
user_48378140

using map and lambda,

n = 3
a = [1, 2, 3, 4, 5, 6]
print(list(set(map(lambda k:sum(k), combinations_with_replacement(a, n)))))
user_11494230
Related