CODEWITHSUNDEEP
python
Tom Carrick
Tom Carrick

Reputation: 6616

Python dict comprehension summing from objects

Given a setup like this:

class Foo():
   state = 'x'
   amount = 1


a = Foo()
b = Foo()
c = Foo()
c.state = 'y'
foos = [a, b, c]

I want to get a dict that has keys = object.state, values = sum(object.amounts of objects with that state). In this case:

{'x': 2, 'y': 1}

I want to do this automatically, so I don't need to know the different possible states in advance.

For sure I could iterate through in some boring manner like this:

my_dict = {}
for foo in foos:
    try:
        my_dict[foo.state] += foo.value
    except (KeyError, TypeError):
        my_dict[foo.state] = foo.value

But that is a bit verbose, and I'm wondering if there's a nicer way to do it, maybe with dict comprehensions or something, but my efforts so far have been in vain.

Upvotes: 3

Views: 523

Answers (3)

Kasravnd
Kasravnd

Reputation: 107347

Dictionary comprehension is not the most optimized approach in this case. Instead you can use collections.defaultdict() like following :

>>> from collections import defaultdict
>>> d = defaultdict(int)
>>> 
>>> for obj in foos:
...     d[obj.state] += obj.amount
... 
>>> d
defaultdict(<type 'int'>, {'y': 1, 'x': 2})

Upvotes: 2

Iron Fist
Iron Fist

Reputation: 10951

How about Counter:

>>> from collections import Counter
>>>
>>> foos = [a,b,c]
>>> 
>>> c = Counter()
>>> for x in foos:
        c[x.state] += x.amount


>>> c
Counter({'x': 2, 'y': 1})

Upvotes: 3

Michael Hunter
Michael Hunter

Reputation: 414

You could use defaultdict.

from collections import defaultdict
my_dict = defaultdict(lambda: 0)
for foo in foos:
    my_dict[foo.type] += foo.value

You could also use setdefault.

my_dict = {}
for foo in foos:
    my_dict.setdefault(foo.type, 0)
    my_dict[foo.type] += foo.value

Upvotes: 1

Related Questions