CODEWITHSUNDEEP
pythonfor-loopdictionary
dassouki
dassouki

Reputation: 6366

Loop through multiple different sized python dictionaries

I am inheriting the following two python dictionaries from another class

>>> print init_treats
{('001', 0): init_treat_001_0, ('001', 3): init_treat_001_3, ('001', 2): init_treat_001_2, ('002', 0): init_treat_002_0, ('002', 1): init_treat_002_1, ('002', 2): init_treat_002_2, ('002', 3): init_treat_002_3, ('001', 1): init_treat_001_1}

>>> print init_untreats
{'002': init_untreat_002, '001': init_untreat_001}

How can I generate the following?

init_treat_001_0 + init_treat_001_1 + init_treat_001_2 + init_treat_001_3 +
init_untreat_001 == 0


init_treat_002_0 + init_treat_002_1 + init_treat_002_2 + init_treat_002_3 + 
init_untreat_002 == 0

Upvotes: 2

Views: 71

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121814

Sort your init_treats keys:

treats = sorted(init_treats)

now you can use itertools.groupby() to group them on the first part of your key:

from itertools import groupby
from operator import itemgetter

for untreat, group in groupby(sorted(init_treats), itemgetter(0)):
    # group is now a sorted iterator of keys with the same first value
    if init_untreat[untreat] + sum(map(init_treats.get, group)) == 0:
        # sum of init_treat_n_m + init_untreat_n is 0

Because this uses sorting, this is a O(NlogN) solution (N being the size of the init_treats dictionary)

You could use a dictionary for a O(N + K) solution (K being the size of the init_untreats dictionary):

sums = init_untreat.copy()
for untreat, id in init_treats:
    sums[untreat] += init_treats[untreat, id]

for untreat, total in sums.items():  # use sums.iteritems() in Python 2
    if total == 0:
        # sum of init_treat_n_m + init_untreat_n is 0

Because K is always smaller than N in your case, asymptotically speaking this is a O(N) algorithm, of course.

Upvotes: 2

Related Questions