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?
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.