Merge dictionary keys when they're almost equal

Viewed 70

I wish to merge some values in a dictionary based on the conditions:

  • dict keys are within +-1 of each other (1 is example, can be a different value)
  • Create new dictionary with averaged key and merged values

A small workable example. Say you have value z and values belonging to that value z called ids:

#Input data
ids = [1,0,2,4,3,5,6,7,8] 
z =   [0,1,1,4,0,1,3,3,1]

#Rewriting in dictionary
dictionary = {}
for item_index, item in enumerate(z):
    if item in dictionary:
        dictionary[item].append(ids[item_index])
    else:
        dictionary[item] = [ids[item_index]]

#Output is now: {0: [1, 3], 1: [0, 2, 5, 8], 4: [4], 3: [6, 7]}

My desired output is based on the fact that keys +- 1 should be merged and averaged:

{0.5: [1,3,0,2,5,8], 3.5: [4,6,7]}

Sorted or not does not matter. Does someone know how to get the desired output in an efficient way? I'm really stuck.

EDIT

The +-1 is an example. I wish to be able to alter the tolerance of merging, and so it should be a variable. Furthermore, the to be merged keys are not always pairs, but can be larger groups

4 Answers

This will work

from more_itertools import consecutive_groups
from itertools import groupby

ids = [1,0,2,4,3,5,6,7,8] 
z =   [0,1,1,4,0,1,3,3,1]

# group by the keys in your dictionary
tmp = sorted(zip(z, ids)) # groupby groups consecutive values: need to sort
groups = groupby(tmp, key = lambda x : x[0])

# now group consecutive keys together
cons_groups = consecutive_groups(groups, ordering = lambda x : x[0])
out_dict = {}

for group in cons_groups:
    result_key = 0
    count = 0
    result = []
    for key, value in group:
        result_key += key
        count += 1
        result.extend(x[1] for x in value)
    out_dict[result_key / count] = result
# {0.5: [1, 3, 0, 2, 5, 8], 3.5: [6, 7, 4]}

The general idea here is:

  1. Group by dictionary keys
  2. Aggregate the above groups into larger groups containing consecutive keys as well.
  3. In the for loop we unravel the above aggregation to get your desired key : value result.

If we add a 2 to your keys like:

ids = [1,0,2,4,3,5,6,7,8,3] 
z =   [0,1,1,4,0,1,3,3,1,2]

the result is now:

{2.0: [1, 3, 0, 2, 5, 8, 3, 6, 7, 4]}

as you expect.

Get the sorted list of keys in reverse order, and create a result variable to store resulting dictionary. Now iterate in loop, and pop a key, if keys list is not empty, pop another key, and check if the difference is positive one, if yes, get the average of the keys and add both the list, else append back the popped keys, but if key list is empty, then just append the current key, loop will terminate.

dictionary={0: [1, 3], 1: [0, 2, 5, 8], 4: [4], 3: [6, 7]}
keys = sorted([key for key in dictionary], reverse=True)

result = {}
while keys:
    currKey = keys.pop()
    if keys:
        nextKey = keys.pop()
        if nextKey-currKey==1:
            newKey = (nextKey+currKey)/2
            result[newKey] = dictionary[currKey]+dictionary[nextKey]
            continue
        else:
            keys.append(nextKey)
            continue
    result[currKey] = dictionary[currKey]

OUTPUT:

{0.5: [1, 3, 0, 2, 5, 8], 3.5: [6, 7, 4]}

UPDATED ANSWER

Above answer was based on the data sample you had, the answer below will work for any number of consecutive keys, logic is still almost the same..

result = {}
while keys:
    currKey = keys.pop()
    if keys:
        tempKeys = [currKey]
        tempValues = dictionary[currKey]
        nextKey = keys.pop()
        while nextKey-currKey == 1:
            tempKeys.append(nextKey)
            tempValues.extend(dictionary[nextKey])
            if not keys:
                break
            currKey = nextKey
            nextKey = keys.pop()
        else:
            keys.append(nextKey)
        result[sum(tempKeys)/len(tempKeys)] = tempValues
        continue

    result[currKey] = dictionary[currKey] 

Who said homemade isn't okay?

It's not too complicated actually. First we group up all the consecutive keys [[0, 1], [3, 4], [8]] Then we find the average of these sublists and store them as keys in a dictionary. For the values of these keys, we iterate through the sublists elements and retrieve each value.

import more_itertools as mit
bad = {0: [1, 3], 1: [0, 2, 5, 8], 4: [4], 3: [6, 7],8:[2]}
grouped = [list(group) for group in mit.consecutive_groups(sorted(list(bad.keys())))]
groupedNew = {sum(lst)/len(lst):[g for num in lst for g in bad[num]] for lst in grouped}

No way that could be one line. (Disclaimer: line count does not mean faster)

import more_itertools as mit
groupedNew = {sum(lst)/len(lst):[g for num in lst for g in {0: [1, 3], 1: [0, 2, 5, 8], 4: [4], 3: [6, 7],8:[2]}[num]] for lst in [list(group) for group in mit.consecutive_groups(sorted(list({0: [1, 3], 1: [0, 2, 5, 8], 4: [4], 3: [6, 7],8:[2]}.keys())))]}`

output

{0.5: [1, 3, 0, 2, 5, 8], 3.5: [6, 7, 4], 8.0: [2]}
from itertools import chain

# sort the keys (d = your dictionary)
sorted_keys = sorted(d.keys())

# get the difference between consecutive (sorted) keys
diff_keys = [next_ - prev for prev, next_ in zip(sorted_keys[1:], sorted_keys)]

# start forming the groups
groups = [[sorted_keys[0]]]

# group index starts from 0
gr = 0

# how much diff is okay to merge?
tolerance = 1

# for each diff..
for j, diff in enumerate(diff_keys, start=1):
    # diff is within the tolerance?
    if diff >= -tolerance:
        # then append to current group
        groups[gr].append(sorted_keys[j])
    else:
        # otherwise start a new group
        groups.append([sorted_keys[j]])
        gr += 1

# merge the groups with a dict-comprehension over groups
{sum(gr) / len(gr): list(chain.from_iterable(map(d.get, gr))) for gr in groups}
Related