Search code examples
pythondictionarylist-comprehensionpython-itertoolsdefaultdict

Python: Lookup keys which correspond to an item in a dictionary of lists


Problem Statement

Given a dictionary of lists,

key_to_list = {
    'one': [1, 3, 5, 7],
    'two': [2, 4, 6, 8],
    'three': [1, 2, 5, 6],
    'four': [2, 5, 7, 8]
}

what is the best way to create a mapping from elements of the lists to their keys?

list_element_to_keys = {
    1: {'one', 'three'},
    2: {'two', 'three', 'four'},
    3: {'one'},
    4: {'two'},
    5: {'one', 'three', 'four'},
    6: {'two', 'three'},
    7: {'one', 'four'},
    8: {'two', 'four'}
}

My Solution

from collections import defaultdict

list_element_to_keys = defaultdict(set)
for key, value in key_to_list.items():
    for item in value:
        list_element_to_keys[item].add(key)

Thoughts

A friend of mine has suggested it may be possible to use a dictionary comprehension, but I keep running into issues because multiple keys have lists that contain some of the same items.

I also think their might be some itertools magic that could help, but I am not positive.

Dict Comprehension

I have found, with the help of a friend a dictionary comprehension that works.

from itertools import chain
list_element_to_keys= { i: set(k for k,v in key_to_list.items() if i in v) for i in set(chain.from_iterable(key_to_list.values())) }

Solution

  • Your solution is fine, it works and defaultdict is an obvious (good) choice for such a problem.

    One thing you could improve is to use six.iteritems(key_to_list) that will make it a bit faster on Python2.