pythonpython-3.xdictionarymergedictionary-comprehension

Summing instead of overriding values in dictionary comprehension


When the keys in a dictionary comprehension is the same, I want their values to be added up. For example,

>>> dct = {-1: 1, 0: 2, 1: 3}
>>> {k**2: v for k, v in dct.items()}
{1: 3, 0: 2}

However, what I want to get in this case is {1: 4, 0: 2}, because both the square of 1 and -1 is 1, and 1 + 3 = 4.

Clearly, I can do it with a for loop, but is there a shorthand?


Solution

  • There isn't a shorthand version, since your comprehension would need to keep track of the current state which isn't doable. Like you said, the answer is a for loop:

    old = {-1: 1, 0: 2, 1:3}
    new = {}
    for k, v in old.items():
        new[k**2] = new.get(k**2, 0) + v
    

    The trick using the dict.get method I saw somewhere in the Python docs. It does the same thing as:

    if k**2 in new:
       new[k**2] += v
    else:
       new[k**2] = v
    

    But this variation uses the get method which returns a default 0 which is added on to the value that will be assigned (when the key doesn't exist). Since it is 0, and the values are numbers being added, 0 has no effect. By contrast, if you needed to get the product, you'd use 1 as the default as starting off with 0 will mean that you never increase the value.

    In addition, the latter, more verbose, method shown above evaluates k**2 twice each cycle which uses up computation. To make it use 1 calculation would require another line of code which in my opinion isn't worth the time when the get method is so much cleaner.