I am a new python user, and I need help about combining list elements under a condition. I have a list like this:
x = [['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
I would like to combine list elements which start with the same letter in a list by summing up the other elements. for example, I'd like to obtain this list for x
:
x = [['a', 30, 120], ['b', 10, 20]]
How can I achieve this ?
A one-liner using itertools.groupby() :
In [45]: lis=[['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
In [46]: lis.sort(key=itemgetter(0)) #sort the list first
In [47]: lis
Out[47]: [['a', 10, 20], ['a', 20, 100], ['b', 10, 20]]
In [49]: [[k]+map(sum,zip(*[x[1:] for x in g])) for k,g in groupby(lis,key=itemgetter(0))]
Out[49]: [['a', 30, 120], ['b', 10, 20]]
A simple solution:
In [23]: lis=[['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
In [24]: ans=[]
In [25]: lis.sort(key=itemgetter(0)) #sort the list according to the first elem
In [26]: lis
Out[26]: [['a', 10, 20], ['a', 20, 100], ['b', 10, 20]]
In [27]: for x in lis:
if ans:
if x[0]==ans[-1][0]: #if the value of the first elem of last element in ans is same as x[0]
ans[-1][1]+=x[1]
ans[-1][2]+=x[2]
else:
ans.append(x)
else:ans.append(x)
....:
In [28]: ans
Out[28]: [['a', 30, 120], ['b', 10, 20]]
Without sorting the list using defaultdict()
:
In [69]: dic=defaultdict(list)
In [70]: for x in lis:
dic[x[0]].append(x[1:])
....:
In [71]: dic
Out[71]: defaultdict(<type 'list'>, {'a': [[10, 20], [20, 100]], 'b': [[10, 20]]})
In [72]: [[k]+map(sum,zip(*i)) for k,i in dic.items()]
Out[72]: [['a', 30, 120], ['b', 10, 20]]
Another approach using dict
and map
:
>>> x = [['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
>>> d = {}
>>> from operator import add
>>> for k, v1, v2 in x:
d[k] = map(add, d.get(k, [0, 0]), [v1, v2])
>>> d
{'a': [30, 120], 'b': [10, 20]}
I'm going to use the answer code for a huge data which include over millons elements. I'd like the reduce the list elements this way.
In such a case you probably don't want to be sorting the data or building a fully copy as you're iterating over it.
The following solution does neither. It can also handle sublists of any length (as long as all lengths are the same):
def add(d, l):
k = l[0] # extract the key
p = d.get(k, None) # see if we already have a partial sum for this key
if p:
d[k] = [x+y for x,y in zip(p, l[1:])] # add to the previous sum
else:
d[k] = l[1:] # create a new sum
return d
x = [['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
result = [[k] + v for k,v in reduce(add, x, {}).items()]
print(result)
Alternatively,
import collections, operator
x = [['a', 10, 20], ['b', 10, 20], ['a', 20, 100]]
d = collections.defaultdict(lambda:[0] * (len(x[0]) - 1))
for el in x:
d[el[0]] = map(operator.add, d[el[0]], el[1:])
result = [[k] + v for k,v in d.items()]
print(result)
This works exactly the same as the first version, but uses defaultdict
and explicit iteration.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.