Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Generalizing adding nested lists

Given 3 nested vectors:

 >>> a
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b
[[10, 20, 30], [40, 50, 60], [70, 80, 90]]
>>> c
[[100, 200, 300], [400, 500, 600], [700, 800, 900]]

I can add these vectors together with a map/sum/zip comprehension like so:

>>> [map(sum,zip(i,j,k)) for i,j,k in zip(a,b,c)]
[[111, 222, 333], [444, 555, 666], [777, 888, 999]]

I've manually expanded this from adding two lists together, but is there a pythonic way to generalize this to handle an arbitrary number of lists?

(Python 2.7 without using external libraries preferred)

like image 831
paulw1128 Avatar asked Nov 09 '16 17:11

paulw1128


2 Answers

Here is a general approach:

from itertools import izip
def multiple_sum(*args):
    return [map(sum, izip(*items)) for items in izip(*args)] 

Demo:

In [13]: multiple_sum(a, b, c)
Out[13]: [[111, 222, 333], [444, 555, 666], [777, 888, 999]]

Note that since in Python 2.7 zip returns a list it's better off using it when you only want to iterate over the results, and instead use itertools.izip that returns an iterator.

Here is another way using itertools.starmap() which is faster than the previouse approach:

def multiple_sum(*args):
    return [map(sum, lst) for lst in starmap(zip, zip(*args))]

Benchmark:

In [32]: %timeit [map(sum, izip(*items)) for items in izip(a, b, c)]
100000 loops, best of 3: 3.93 µs per loop

In [33]: %timeit [map(sum, lst) for lst in starmap(zip, zip(a, b , c))]
100000 loops, best of 3: 3.01 µs per loop
like image 152
Mazdak Avatar answered Nov 02 '22 20:11

Mazdak


This should work:

lists = [a,b,c]
[map(sum,zip(*zipped_lists)) for zipped_lists in zip(*lists)]
like image 26
lucasnadalutti Avatar answered Nov 02 '22 20:11

lucasnadalutti