Python: adding items to a nested list if the first two items are the same

Question

I have a nested list:

 a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

If the first two elements of the auxiliary lists are the same, I want to add their corresponding 3, 4, and 5 elements, keeping the first two elements. For the above case, the result should be:

a = [[1,2,13,12,13],[1,3,4,5,7],[3,4,2,3,4]]

Is there any way to do this?

Thanks.

Answer 1

You can use the dictionary to categorize sub-lists based on the first two elements, then iterate over the elements and calculate the sum:

>>> d = {}
>>> for sub in a:
...    d.setdefault(tuple(sub[:2]),[]).append(sub[2:])
... 

>>> 
>>> [k+tuple(map(sum, zip(*v))) for k,v in d.items()]
[(1, 2, 13, 12, 13), (1, 3, 4, 5, 7), (3, 4, 2, 3, 4)]

Answer 2

Pandas Single Line

How about pandas single line.

Import and data:

>>> import pandas as pd
>>> a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

Pandas in action:

>>> pd.DataFrame(a).groupby([0, 1]).sum().reset_index().values.tolist()
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

Step by step

Create a data block:

>>> df = pd.DataFrame(a)

   0  1  2  3  4
0  1  2  3  4  5
1  1  2  5  6  7
2  1  2  5  2  1
3  1  3  4  5  7
4  3  4  1  2  3
5  3  4  1  1  1

:

>>> df2 = df.groupby([0, 1]).sum()
>>> df2

      2   3   4
0 1            
1 2  13  12  13
  3   4   5   7
3 4   2   3   4

groupby 0 1. GroupBy. "" - , . . , , :

MultiIndex:

MultiIndex(levels=[[1, 3], [2, 3, 4]],
           labels=[[0, 0, 1], [0, 1, 2]],
           names=[0, 1])

0 1. - 1 3, - 2, 3 4. 1 . .

:

>> df3 = df2.reset_index()

   0  1   2   3   4
0  1  2  13  12  13
1  1  3   4   5   7
2  3  4   2   3   4

:

>>> df3.values.tolist()
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

Answer 3

, :

a=[[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]

#create empty dict
d = {}

# function for adding list b to lists a
def addlist(a,b): return [a[i]+b[i] for i in range(len(a))]  

# now iter through each item and add to existing or create a new record in dictionary
for i in a:
    d[tuple(i[:2])] = addlist(d.setdefault(tuple(i[:2]),[0,0,0]),i[2:])

#the wanted output would be
[list(k)+v for k,v in d.items()]

Answer 4

itertools.groupby.

Here's how to do it with a nested list comprehension.

from itertools import groupby
from operator import itemgetter

a = [[1,2,3,4,5],[1,2,5,6,7],[1,2,5,2,1],[1,3,4,5,7],[3,4,1,2,3],[3,4,1,1,1]]
expected = [[1,2,13,12,13],[1,3,4,5,7],[3,4,2,3,4]]
print(expected)

a = [list(k) + [sum(t) for t in zip(*[u[2:] for u in g])]
    for k, g in groupby(a, itemgetter(0, 1))]
print(a)

Exit

[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]
[[1, 2, 13, 12, 13], [1, 3, 4, 5, 7], [3, 4, 2, 3, 4]]

Note that this assumes that the list aalready has sub-lists that start with the same two adjacent elements. If it is not, then it should be sorted, for example

a.sort(key=itemgetter(0, 1))

before running the above code.

Here essentially the same algorithm is broken down for easier reading and analysis.

keyfunc = lambda seq: seq[:2]

a.sort(key=keyfunc)
new_a = []
for k, g in groupby(a, key=keyfunc):
    tails = [u[2:] for u in g]
    sums = [sum(t) for t in zip(*tails)]
    new_a.append(k + sums)

print(new_a)