May 27, 2022

Invert a Complex Dictionary in Python

Have you ever need to invert a dictionary with value being a list in Python?

Here is what I mean. We need to convert the following dict:

{ "apple"      : [ "green", "red" ], 

  "watermelon" : [ "green" ], 

  "strawberry" : [ "red" ], 

  "lemon"      : [ "green", "yellow" ] }

To a new dict as below:

{ 'green'  : ['apple', 'watermelon', 'lemon'], 

  'red'    : ['apple', 'strawberry'], 

  'yellow' : ['lemon'] } 


Here's my solution which using the defaultdict from collections:

>>> from collections import defaultdict

>>> a_dict = { "apple" : [ "green", "red" ], "watermelon" : [ "green" ], "strawberry" : [ "red" ], "lemon" : [ "green", "yellow" ] }

>>> b_dict = defaultdict(list)

>>> for k,v in a_dict.items():

...    for k1 in v:

...        b_dict[k1].append(k)

...

>>> b_dict

defaultdict(<class 'list'>, {'green': ['apple', 'watermelon', 'lemon'], 'red': ['apple', 'strawberry'], 'yellow': ['lemon']})


There is another way that simplies what we have above:

>>> from collections import defaultdict

>>> a_dict = { "apple" : [ "green", "red" ], "watermelon" : [ "green" ], "strawberry" : [ "red" ], "lemon" : [ "green", "yellow" ] }

>>> b_dict = defaultdict(list)

>>> { b_dict[k1].append(k) for k,v in a_dict.items() for k1 in v }

>>> b_dict


Alright, hope this helps!!