Pythonic way to add a list of vectors

Question

I am trying to create a method (sum) that takes a variable number of vectors and adds them in. For educational purposes, I have written my own Vector class, and the underlying data is stored in an instance variable named data.

My code for the @classmethod sum works (for each of the vectors passed in, loop through each element in the data variable and add it to a result list), but it seems non-Pythonic, and wondering if there is a better way?

class Vector(object):
    def __init__(self, data):
        self.data = data

    @classmethod
    def sum(cls, *args):
        result = [0 for _ in range(len(args[0].data))]
        for v in args:
            if len(v.data) != len(result): raise
            for i, element in enumerate(v.data):
                result[i] += element

        return cls(result)

Answer 1

itertools.izip_longest may come very handy in your situation:

a = [1, 2, 3, 4]
b = [1, 2, 3, 4, 5, 6]
c = [1, 2]

lists = (a, b, c)

result = [sum(el) for el in itertools.izip_longest(*lists, fillvalue=0)]

And here you got what you wanted:

>>> result
[3, 6, 6, 8, 5, 6]

What it does is simply zips up your lists together, by filling empty value with 0. e.g. izip_longest(a, b) would be [(1, 1), (2, 2), (3, 0), (4, 0)]. Then just sums up all the values in each tuple element of the intermediate list.

So here you go step by step:

>>> lists
([1, 2, 3, 4], [1, 2, 3, 4, 5, 6], [1, 2])
>>> list(itertools.izip_longest(*lists, fillvalue=0))
[(1, 1, 1), (2, 2, 2), (3, 3, 0), (4, 4, 0), (0, 5, 0), (0, 6, 0)]

So if you run a list comprehension, summing up all sub-elements, you get your result.

Answer 2

Another thing that you could do (and that might be more "pythonic") would be to implement the __add__ magic method, so you can use + and sum directly on vectors.

class Vector(object):
    def __init__(self, data):
        self.data = data

    def __add__(self, other):
        if isinstance(other, Vector):
            return Vector([s + o for s, o in zip(self.data, other.data)])
        if isinstance(other, int):
            return Vector([s + other for s in self.data])
        raise TypeError("can not add %s to vector" % other)

    def __radd__(self, other):
        return self.__add__(other)

    def __repr__(self):
        return "Vector(%r)" % self.data

Here, I also implemented addition of Vector and int, adding the number on each of the Vector's data elements, and the "reverse addition" __radd__, to make sum work properly.

Example:

>>> v1 = Vector([1,2,3])
>>> v2 = Vector([4,5,6])
>>> v3 = Vector([7,8,9])
>>> v1 + v2 + v3
Vector([12, 15, 18])
>>> sum([v1,v2,v3])
Vector([12, 15, 18])