CODEWITHSUNDEEP
pythonlist
Black
Black

Reputation: 4644

Python: Add lists together

I'm trying to add the last four elements of k lists together element-wise given these lists have the same i[0] value.

lst = [[1,1,1],[2,1,2],[2,2,1],[3,4,1],[3,4,5],[3,1,0]]

Output = [[1,1,1],[2,3,3],[3,9,6]]

I know I should use:

from operator import add
map(add, list1, list2, ..., listk)

My issue if that there could be k-lists with the same i[0] value so I am not sure how to incorporate this element. Is there an efficient way to do this?

Upvotes: 0

Views: 205

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121524

Use itertools.groupby():

from itertools import groupby
from operator import itemgetter

[[k] + [sum(col) for col in zip(*g)[1:]] for k, g in groupby(lst, itemgetter(0))]

The zip(*iterable) call allows us to sum the group values per column; the [1:] slice lets us ignore the first column.

Demo:

>>> from itertools import groupby
>>> from operator import itemgetter
>>> lst = [[1, 1, 1], [2, 1, 2], [2, 2, 1], [3, 4, 1], [3, 4, 5], [3, 1, 0]]
>>> [[k] + [sum(col) for col in zip(*g)[1:]] for k, g in groupby(lst, itemgetter(0))]
[[1, 1, 1], [2, 3, 3], [3, 9, 6]]

Upvotes: 3

Related Questions