CODEWITHSUNDEEP
python
JRR
JRR

Reputation: 6152

efficient way to insert elements into a dictionary

Is there a simpler way to write this code?

def insert(k, v)  # v can either be a single number or a list of numbers, and k may or may not exist in self.dict
  if k in self.dict:
    if isinstance(v, list):
      for e in v:
        self.dict[k].add(e)
    else:
      self.dict[k].add(v)
  else:
    self.dict[k] = {}
    if isinstance(v, list):
      for e in v:
        self.dict[k].add(e)
    else:
      self.dict[k].add(v)

Upvotes: 2

Views: 231

Answers (4)

Todd
Todd

Reputation: 5405

Here's an exercise in the problem showing use of exception handling and limited recursion - this may be the most efficient way to do it. Although it 'looks' like the code is larger than other solutions, it compiles down into fewer instructions than the others:

>>> class Foo:
...     def __init__(self):
...         self.dict = {}
... 
...     def insert(self, k, v):
...         try:
...             self.dict[k].extend(v)  # <-- majority of accesses 
...                                     #     only execute this one line.
...         except KeyError:
...             self.dict[k] = []
...             self.insert(k, v)
...         except TypeError:
...             self.dict[k].extend([v])
...             
>>> f = Foo()
>>> f.insert('bob', 1)
>>> f.dict
{'bob': [1]}
>>> f.insert('bob', [3, 4])
>>> f.dict
{'bob': [1, 3, 4]}

I believe this will have the fastest performance assuming most of your data is lists of numbers with only occasional single digits. And that the keys don't vary wildly. The single line self.dict[k].extend(v) should be executed without exceptions the majority of accesses to the object. It incurs no extra cost checking types and testing conditions. Even if single digits were frequent, branching to add them incurs about the same cost as type checking - except this implementation isn't forced to type check every time something is added.

Upvotes: 0

abc
abc

Reputation: 11929

You could use a defaultdict and simplify the code to

def insert(self, k, v): 
   if isinstance(v, list):
      self.dict[k]+=v
   else:
      self.dict[k].append(v)

Note that in this case the dictionary would need to be defined as:

self.dict = defaultdict(list)

Or, if you want to associate a set to your keys you could just change it to defaultdict(set) and use:

def insert(self, k, v): 
   if isinstance(v, list):
      self.dict[k]|=set(v)
   else:
      self.dict[k].add(v)

Upvotes: 2

user2390182
user2390182

Reputation: 73480

Without any changes to the types you use, you can also just do the following in one line:

def insert(k, v): 
    d.setdefault(k, []).extend(v if isinstance(v, list) else [v])

Or, when using sets:

def insert(k, v): 
    d.setdefault(k, set()).update(v if isinstance(v, list) else [v])

Upvotes: 0

Masklinn
Masklinn

Reputation: 42492

def insert(k, v):
    entries = self.dict.setdefault(k, set())
    if isinstance(v, list):
        entries.update(v)
    else:
        entries.add(v)

Upvotes: 0

Related Questions