Bennox75
Bennox75

Reputation: 33

Merge lists in a dictionary that have a common element

I would like to merge both keys and values of a dictionary if the values (lists) that share at least one element.

The input would be :

dico = {"a" : [1,2,3], "b":[9,2,89], "c":[3,12,530],"d":[34,42],"e":[34,6]}

And I would like the out to be something like this :

{"a,b,c" : [1,2,3,9,89,12,530], "d,e": [34,42,6] }

None of what I tried worked... Do you think this is possible?

Upvotes: 1

Views: 278

Answers (1)

tobias_k
tobias_k

Reputation: 82899

You an use a Union-Find aka Disjoin Set approach. First, you need two functions: union and find. I usually keep those lying around somewhere in case I need them.

def find(x):
    l = leaders[x]
    if l is not None:
        l = find(l)
        leaders[x] = l
        return l
    return x

def union(x, y):
    lx, ly = find(x), find(y)
    if lx != ly:
        leaders[lx] = ly

Now, you can use those to determine one "leader" for each element in the lists...

dico = {"a" : [1,2,3], "b":[9,2,89], "c":[3,12,530],"d":[34,42],"e":[34,6]}
leaders = collections.defaultdict(lambda: None)

for val in dico.values():
    for other in val[1:]:
        union(val[0], other)        

... and then group elements that have the same "leader" into groups.

groups = collections.defaultdict(set)
for x in leaders:
    groups[find(x)].add(x)

Now, also group the keys by the leaders of their first elements:

keys = collections.defaultdict(list)
for key in dico:
    keys[find(dico[key][0])].append(key)

And finally assemble the result.

result = {','.join(ks): groups[leader] for (leader, ks) in keys.items()}
# {'d,e': {42, 34, 6}, 'c,a,b': {1, 2, 3, 9, 12, 530, 89}}

Note that this is using sets, though, instead of lists. If you need to retain the original order, just group the keys and then chair their respective lists together.

Upvotes: 1

Related Questions