Add second elements of a list based on equality condition on Python

Viewed 88

Let's say I have a list like this :

[[0.5, 5281],
 [0.7, 6597],
 [0.7, 6716],
 [0.7, 6902],
 [0.7, 5704]]

I want to sum the elements that have the same 1st element and get something like :

[[0.5, 5281],
 [0.7, the result of 6597+6716+..+5704]].

IS there an easy way to do that in Python?

8 Answers

What about using pandas? (I suppose you read your database with pandas anyway)

import pandas as pd

original_list = [
    [0.5, 5281],
    [0.7, 6597],
    [0.7, 6716],
    [0.7, 6902],
    [0.7, 5704]]

df = pd.DataFrame(original_list, columns=['col1', 'col2'])
df_out = df.groupby('col1').sum()
print(df_out)

Result:

       col2
col1       
0.5    5281
0.7   25919

This I find fairly simple. You can create an empty dict and append it's keys and values. If first part of the sublist already exists as a key, add the second part of the sublist to it's value:

l = [[0.5, 5281], [0.7, 6597], [0.7, 6716], [0.7, 6902], [0.7, 5704]]
d = {}

for item in l:
   if item[0] in d.keys():
      d[item[0]] += item[1]
   else:
      d[item[0]] = item[1]

Result:

>>> print(d)
{0.5: 5281, 0.7: 25919}

>>> print(list(map(list, d.items())))
[[0.5, 5281], [0.7, 25919]]

Another solution that doesn't use pandas (but I wouldn't say it is simpler, perhaps, even the opposite). You can use groupby from itertools to group the items (but first you need to sort the list (unless you already receive it sorted (in which case the sorting wouldn't consume that many resources anyways)) because otherwise it may not group all of the same items together, just the ones that follow each other). Then just append the group and the sum to a list:

import itertools
import operator


original_list = [
    [0.5, 5281],
    [0.7, 6597],
    [0.7, 6716],
    [0.7, 6902],
    [0.7, 5704]]


key = operator.itemgetter(0)
out_list = []
sorted_list = sorted(original_list, key=key)
for group, items in itertools.groupby(sorted_list, key=key):
    out_list.append([group, sum(x[1] for x in items)])

print(out_list)

simple and without any module, do it like this-

ll = [[0.5, 5281],
      [0.7, 6597],
      [0.7, 6716],
      [0.7, 6902],
      [0.7, 5704]]

res = {}
for l in ll:
    if l[0] in res.keys():
        res[l[0]] += [l[1]]
    else:
        res[l[0]] = [l[1]]

for k, v in res.items():
    res[k] = sum(v)

print(res)

Try built-in itertools:

import itertools
import operator

DATA = [[0.5, 5281],
        [0.7, 6597],
        [0.7, 6716],
        [0.7, 6902],
        [0.7, 5704]]

if __name__ == "__main__":
    result = [
        [key, sum(value for _, value in original_pairs)]
        for key, original_pairs in itertools.groupby(DATA, operator.itemgetter(0))
    ]
    print(result)

Using numpy as follows,

import numpy as np

def func_(input_):
    input_array = np.array(input_)
    result = []

    for unique_elem in np.unique(input_array[:,0]):
        indices = np.where(input_array[:,0] == unique_elem)
        result.append([unique_elem, np.sum(input_array[:,1][indices])])

    return result

original_list = [[0.5, 5281],
    [0.7, 6597],
    [0.7, 6716],
    [0.7, 6902],
    [0.7, 5704]]

print(func_(input_))

Output:

[[0.5, 5281.0], [0.7, 25919.0]]

This can still be improved

Since this is the same as the classic "wordcount" problem, you could approach this in a MapReduce style using functools.reduce:

from operator import itemgetter
from functools import reduce

myList=[[0.5, 5281],
 [0.7, 6597],
 [0.7, 6716],
 [0.7, 6902],
 [0.7, 5704]]

# sort data by first coordinate (key)
# in Hadoop this is done by the mapper 
sorted(myList, key=itemgetter(0))
# Out: [[0.5, 5281], [0.7, 6597], [0.7, 6716], [0.7, 6902], [0.7, 5704]]

# reducer function that adds up items with the same key 
# this works because input list is sorted by key
def myReducer(a,b):
    if a!=[] and b[0] == a[-1][0]:
        c = a.pop()
        return a + [[c[0],c[1]+b[1]]]
    else:
        return a + [b]

# putting it all together
reduce(myReducer, sorted(myList, key=itemgetter(0)), [])
# Out: [[0.5, 5281], [0.7, 25919]]

This approach is also suitable for parallelization.

Yeah sure its doable! If you want to do it from scratch without any imported libraries:

A = [
    [0.5, 5281],
    [0.7, 6597],
    [0.7, 6716],
    [0.7, 6902],
    [0.7, 5704]
    ]


table_of_first_entries = []
result = []

for input_pair in A:
    if input_pair[0] not in table_of_first_entries:
        table_of_first_entries.append(input_pair[0])

for first_entry in table_of_first_entries:
    sum = 0
    for input_pair in A:
        if input_pair[0] == first_entry:
            sum = sum + input_pair[1]
    
    result.append([first_entry,sum])


print(result)

and the output you get would be

[[0.5, 5281], [0.7, 25919]]
Related