ignore my earlier comment... sorry I misunderstood it. 

Here, instead of this if else, you can use defaultdict.