Python Dictionary: Replace values based on key occurrence in other values

Question

My question might involve comprehension if I'm not mistaken.
So I have a dictionary as:

data = {
 1: [2, 4], 
 2: [3], 
 3: [1, 2, 6], 
 4: [1, 3, 6], 
 6: [3, 4]
}  

I wish to rearrange the dictionary values (not keys) such that with every occurrence of a "key" in any of the other values, the current key's values should now contain the "key" in who's value it originally occurred.
For example the dictionary should now look like this:

new_data = {
 1: [3, 4], 
 2: [1, 3], 
 3: [2, 4, 6], 
 4: [1, 6], 
 6: [3, 4]
}  

I know for sure that I must use another empty dictionary to solve this, and I must use looping. I have a starting point, but it's not efficient:

for pk,pv in zip(data.keys(), data.values()):
    #print(pk,pv)
    for pvs in pv:
        #print(pvs)
        if pk == pvs:
            new_data.setdefault(pk, [])
            new_data[pk].append(pvs)

print(new_data)  

Any help is appreciated. Thanks in advance :)
(Using Ubuntu 14.04 32-Bit VM and Python 2.7)

Answer 1

Perhaps this is not the most efficient way to do this, but I would argue it's one of the nicest to read :).

One thing I should say before, if it's possible to change the initial dictionary values to set instead of list , that could result in some performance gains since set membership checks can be faster than lists s.

Here goes:

def present_in(k, d):
    return {idx for idx in d if k in d[idx]}

new_data = {k: present_in(k, data) for k in data}

Basically the present_in function creates a set of values corresponding to keys where k is present as a member. We simply then use this to create the new_data dictionary using a dictionary comprehension. Note again the values of this dictionary is a set, instead of a list. You can easily wrap them in list() if you do need lists.

EDIT:

If you would like new_data to have keys that may be missing as keys in data but present as one of the values, you can change the code to create a set of all possible keys first:

all_keys = {v for vs in data.values() for v in vs}

def present_in(k, d):
    return {idx for idx in d if k in d[idx]}

new_data = {k: present_in(k, data) for k in all_keys}

Answer 2

{k:[i for i,j in data.items() if k in j] for k,_ in data.items()}
#Output:
#{1: [3, 4], 2: [1, 3], 3: [2, 4, 6], 4: [1, 6], 6: [3, 4]}

The below part essentially checks if the key(k from outer dict) is present in any of the values in dict, if present, it'll get the corresponding key into the list

[i for i,j in data.items() if k in j]