Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Sum of elements in a list of lists with varying lengths in Python

Tags:

python

list

I am trying to calculate the sum of elements in a list of lists. I have no trouble in computing the sum if the lists within the main list all have the same size, like below:

a = [[4], [8], [15]]
total = [sum(i) for i in zip(*a)]

result:

total = [27]     #(4 + 8 + 15) = 27, GOOD!!!

however there is a chance that I may have a list with a different size within the main list, such as:

a = [[3], [4, 6], [10]]

expected result:

total = [17, 19]     #(3 + 4 + 10) = 17, (3 + 6 + 10) = 19

I got stuck here, obviously my solution for the lists of equal sizes does not work. What would be the best way to get the result the way I defined? My intuition was figuring out the list with the maximum length, then expanding other lists to that length with adding zeros, and finally computing the sums separately. This sounds like an ugly solution though, and I wonder if there is there a quick and more elegant way to do this.

Thanks!

EDIT: Should have explained it better. I got confused a bit too... Here below are better examples:

The number of elements in the lists within the list a never exceeds 2. Examples:

a = [[1], [10], [5]] #Expected result: [16] (1+10+5)

a = [[1, 10], [3], [4]] #Expected result: [8, 17] (1+3+4, 10+3+4)

a = [[1, 10], [3], [2, 8]] #Expected result: [6, 12, 15, 21] (1+3+2, 1+3+8, 10+3+2, 10+3+8)

EDIT2: Accepted answer computes the correct results independent of the list sizes.

like image 768
marillion Avatar asked Jan 28 '26 07:01

marillion


1 Answers

Wild guess: you want every possible sum, i.e. the sums you get from taking every possible selection of elements from the sublists?

>>> from itertools import product
>>> a = [[4], [8], [15]]
>>> [sum(p) for p in product(*a)]
[27]
>>> a = [[3], [4, 6], [10]]
>>> [sum(p) for p in product(*a)]
[17, 19]

One way to check this interpretation is to see whether you like the answer it gives for the test in the comments:

>>> a = [[1,2], [3,4,5], [6,7,8,9]] # Tim Pietzcker's example
>>> [sum(p) for p in product(*a)]
[10, 11, 12, 13, 11, 12, 13, 14, 12, 13, 14, 15, 11, 12, 13, 14, 12, 13, 14, 15, 13, 14, 15, 16]
like image 80
DSM Avatar answered Jan 29 '26 20:01

DSM