JHixson JHixson - 11 months ago 42
Python Question

Pythonic way to add values to a set within a dictionary

Lets say i have a dictionary of sets:

d = {"foo":{1,2,3},

Now lets say I want to add the value
to the set found within the key
. This would be easy:


but what if we were unsure of the key already existing? It doesn't feel very pythonic to check beforehand:

if "baz" in dict:
d["baz"] = {7}

I tried to be clever and do something like

d["baz"] = set(d["baz"]).add(7)

but then you just get a
trying to access a bad key in the

Am i missing something, or do I need to just bite the bullet and look before I leap? I would understand if that were the case, it would just be neat if there were a simple way to say "Add this value to the set found at this location, or if there isn't a set at that location, make one, and then put it in.

Answer Source

Use defaultdict

>>> from collections import defaultdict
>>> d = defaultdict(set)
>>> d
defaultdict(<class 'set'>, {})
>>> d['foo'].add(1)
>>> d['foo'].add(2)
>>> d
defaultdict(<class 'set'>, {'foo': {1, 2}})
>>> d['bar'].add(3)
>>> d['bar'].add(4)
>>> d
defaultdict(<class 'set'>, {'foo': {1, 2}, 'bar': {3, 4}})

Also, if you must use plain dict, you can use the .setdefault method:

>>> d2 = {}
>>> d2.setdefault('foo',set()).add(1)
>>> d2.setdefault('foo',set()).add(2)
>>> d2
{'foo': {1, 2}}
>>> d2.setdefault('bar',set()).add(3)
>>> d2.setdefault('bar',set()).add(4)
>>> d2
{'foo': {1, 2}, 'bar': {3, 4}}

Edit to add time comparisons

You should note that using defaultdict is faster:

>>> setup = "gen = ((letter,k) for letter in 'abcdefghijklmnopqrstuvwxyx'  for k in range(100)); d = {}"
>>> s = """for l,n in gen:
...     d.setdefault(l,set()).add(n)"""
>>> setup2 = "from collections import defaultdict; gen = ((letter,k) for letter in 'abcdefghijklmnopqrstuvwxyx'  for k in range(100)); d = defaultdict(set)"
>>> s2 = """for l,n in gen:
...     d[l]=n"""
>>> import timeit
>>> timeit.timeit(stmt=s, setup=setup, number=10000)
>>> timeit.timeit(stmt=s2, setup=setup2, number=10000)