paulw1128 paulw1128 - 23 days ago 12
Python Question

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)

Answer

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
Comments