tags:

views:

2010

answers:

4

I'm pretty sure there should be a more Pythonic way of doing this - but I can't think of one: How can I merge a two-dimensional list into a one-dimensional list? Sort of like zip/map but with more than two iterators.

Example - I have the following list:

array = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

I want to have

result = [12, 15, 18] # [1+4+7, 2+5+8, 3+6+9]

So far what I've come up with is:

def add_list(array):
    number_items = len(array[0])
    result = [0] * number_items
    for index in range(number_items):
        for line in array:
            result[index] += line[index]
    return result

But that doesn't look very elegant/Pythonic to me. Aside from not checking if all the "lines" in the 2D array are of the same length, can be added to each other, etc. What would be a better way to do it?

+25  A: 
[sum(a) for a in zip(*array)]
Ned Batchelder
+17  A: 

[sum(value) for value in zip(*array)] is pretty standard.

This might help you understand it:

In [1]: array=[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

In [2]: array
Out[2]: [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

In [3]: *array
------------------------------------------------------------
   File "<ipython console>", line 1
     *array
     ^
<type 'exceptions.SyntaxError'>: invalid syntax

The unary star is not an operator by itself. It unwraps array elements into arguments into function calls.

In [4]: zip(*array)
Out[4]: [(1, 4, 7), (2, 5, 8), (3, 6, 9)]

zip() is a built-in function

In [5]: zip(*array)[0]
Out[5]: (1, 4, 7)

each element for the list returned by zip is a set of numbers you want.

In [6]: sum(zip(*array)[0])
Out[6]: 12

In [7]: [sum(values) for values in zip(*array)]
Out[7]: [12, 15, 18]
Charles Merriam
Nicely answered, Charles.
ΤΖΩΤΖΙΟΥ
Thanks for the explanation of the * operator! It's a good touch.
ojrac
+2  A: 

If you're doing a lot of this kind of thing, you want to learn about scipy.

>>> import scipy
>>> sum(scipy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
array([12, 15, 18])

All array sizes are checked for you automatically. The sum is done in pure C, so it's very fast. scipy arrays are also very memory efficient.

The drawback is you're dependent on a fairly complex third-party module. But that's a very good tradeoff for many purposes.

fivebells
+4  A: 

An alternative way:

map(sum, zip(*array))
ngn