CODEWITHSUNDEEP
python
Liao Zhuodi
Liao Zhuodi

Reputation: 3293

One-liner to create dictionary of lists

how to do this list/dict comprehension turn this [("a", 1), ("b", 2), ("a", 3)]

into this

{
  "a": [1, 3],
  "b": [2]
}

I know how to do it in a for loop, can I use just one line to do the work?

Upvotes: 5

Views: 822

Answers (2)

user3483203
user3483203

Reputation: 51155

Personally, I think using a loop is a simple option here, but for the sake of demonstration, you can define a custom class here that inherits from collections.MutableMapping to accomplish your goal.

import collections

class AutoListDict(collections.MutableMapping):

  def __init__(self, values=()):            
      self.inner = collections.defaultdict(list)
      for k, v in values:
          self[k].append(v)

  def __setitem__(self, k, v): self.inner[k] = v
  def __getitem__(self, k):return self.inner[k]
  def __delitem__(self, k_): del self.inner[k]
  def __iter__(self): return iter(self.inner)
  def __len__(self): return len(self.inner)
  def __repr__(self): return str(dict(self.inner))

This initializes a defaultdict from your input values, so instead of overwriting a key it appends to it. In action, you can simply call the class constructor on your array of tuples:

>>> arr = [('a', 1), ('b', 2), ('a', 3)]
>>> AutoListDict(arr)
{'a': [1, 3], 'b': [2]}

Or if you want something more similar to a dictionary comprehension:

>>> dct = AutoListDict((k, v) for k, v in arr)
>>> dct
{'a': [1, 3], 'b': [2]}

Upvotes: 2

RoadRunner
RoadRunner

Reputation: 26315

A simple approach is to use a simple collections.defaultdict() of lists:

from collections import defaultdict

lst = [("a", 1), ("b", 2), ("a", 3)]

items = defaultdict(list)
for k, v in lst:
    items[k].append(v)

print(items)

Which creates:

defaultdict(<class 'list'>, {'a': [1, 3], 'b': [2]})

Note: If you want the final result to be a normal dictionary, you can just wrap dict().

If you really want a one liner, you could use itertools.groupby() in a dict comprehension:

>>> from itertools import groupby
>>> from operator import itemgetter
>>> lst = [("a", 1), ("b", 2), ("a", 3)]
>>> {k: list(map(itemgetter(1), g)) for k, g in groupby(sorted(lst, key=itemgetter(0)), key=itemgetter(0))}
{'a': [1, 3], 'b': [2]}

Which can also be written more cleanly as:

{k: [x[1] for x in g] for k, g in groupby(sorted(lst, key=itemgetter(0)), key=itemgetter(0))}

The above solution has O(NlogN) complexity due to sorting, a requirement if you want to group similar items together. This is less efficient than The first defaultdict solution, which is O(N), since you only need to iterate the list once. The first solution would be more preferable, since its easier to read, efficient and maintainable.

Upvotes: 5

Related Questions