Search code examples
pythonnested-lists

How to merge the elements of nested list with changing inside elements?


I have a nested list. Two of inside elements can be a same, one is different. I need to change that list in some way that better to demonstrate.

before = [['a','a', 1],['a','a', 2], ['a','b', 6], ['a','b', 0], ['a','c',5]
#after some changes
after = [['a','a', 3], ['a','b', 6], ['a','c',5]]

So i need to summarize the "i[2] elements for i in before" if i[0],i[1] are same. That formulation is already looking like comprehension, but i see no ways to do it in this way.

I did it like that. But i sure it's not most elegant solution.

before = [['a','a', 1],['a','a', 2], ['a','b', 6], ['a','b', 0], ['a','c',5]]
#creating a list of unique combinations, actually it's already created and used for another aims too
after = set([(i[0],i[1],0) for i in before])
after = [list(i) for i in after]
# output is [['a', 'c', 0], ['a', 'b', 0], ['a', 'a', 0]]

for k,i in enumerate(after):
    #starting the loop from list of uniq_combinations
    q = 0 #for element we need to summarize, reseting var for each iteration
    for j in before:
            if (i[0],i[1]) == (j[0],j[1]):
                q = j[2] #the third element for each nested list
                after[k][2] +=q #summarizing the result while looking for same combination
print(after)
#output have correct result [['a', 'a', 3], ['a', 'b', 6], ['a', 'c', 5]]

I've looked in itertools lib, but can't find suitable method


Solution

  • You can use itertools.groupby, as long as elements with the same key are always adjacent.

    from itertools import groupby
    l = [['a','a', 1],['a','a', 2], ['a','b', 6], ['a','b', 0], ['a','c',5]]
    res = [[*k, sum(o[2] for o in g)] for k, g in groupby(l, lambda o:o[:2])]
    

    A defaultdict can also be used to store the sum for each key.

    from collections import defaultdict
    d = defaultdict(int)
    for *k, v in l: d[tuple(k)] += v
    res = [[*k, v] for k, v in d.items()]