Python: select n of each type from ordered dict

Question

I have a for loop that contains multiple if -conditions. The speed of this function is not critical for my application at the moment but I'd like to learn and have a few tips how to optimize my code to make it run faster. In my loop, I get first 5 A's, 5 B's and first 2 C's of each type C_1 and C_2. The code is working correctly, I get a wanted output but I'd like to make it as efficient as possible as I'm not sure how big amount of data it will be given at some point later. Here's the code:

from collections import OrderedDict


odl = [OrderedDict([('id', '1'), ('date', '2022-01-08'), ('time', '15:59:00'), ('type', 'A')]), OrderedDict([('id', '8'), ('date', '2022-01-08'), ('time', '14:59:00'), ('type', 'A')]), OrderedDict([('id', '2'), ('date', '2022-01-09'), ('time', '11:59:00'), ('type', 'A')]),  OrderedDict([('id', '3'), ('date', '2022-01-08'), ('time', '12:59:00'),  ('type', 'B')]), OrderedDict([('id', '9'), ('date', '2022-01-09'), ('time', '17:59:00'),  ('type', 'B')]), OrderedDict([('id', '4'), ('date', '2022-01-09'), ('time', '18:59:00'), ('type', 'B')]),
       OrderedDict([('id', '5'), ('date', '2022-01-08'), ('time', '09:59:00'), ('type', 'C'), ('C_TYPE', 'C_1')]), OrderedDict([('id', '6'), ('date', '2022-01-09'), ('time', '10:59:00'), ('type', 'C'), ('C_TYPE', 'C_2')]), OrderedDict([('id', '7'), ('date', '2022-01-07'), ('time', '16:59:00'),  ('type', 'A')])]

sorted_list = sorted(odl, key=lambda item: (item['date'], item['time']))
sl_list = []
type_a = 0
type_b = 0
type_c_1 = 0
type_c_2 = 0
for s in sorted_list:
    if s['type'] == 'A':
        if type_a < 5:
            sl_list.append(s)
            type_a += 1
    if s['type'] == 'B':
        if type_b < 5:
            sl_list.append(s)
            type_b += 1
    if s['type'] == 'C':
        if s['C_TYPE'] == 'C_1':
            if type_c_1 < 2:
                sl_list.append(s)
                type_c_1 += 1
        if s['C_TYPE'] == 'C_2':
            if type_c_2 < 2:
                sl_list.append(s)
                type_c_2 += 1

Example output:

[OrderedDict([('id', '7'), ('date', '2022-01-07'), ('time', '16:59:00'), ('type', 'A')]), ...]

Are there some obvious performance issues I could face with bigger amounts of data? I'd appreciate every tip to make this more efficient!

Answer 1

Use a collection rather than spawning lots of similarly named variables. Anytime you find yourself creating a bunch of numbered/lettered variables, stop and figure out a way to put that information in a collection. If your data is smart, your code can often be simple.

Applying limits when collecting data by type. You appear to want to control the number of dicts of each type that you collect. Rather than weaving those limits into your algorithm, define a data structure holding the maximums. And then while collecting the data, use a copy of that collection to keep track of how many of each type are still available to be selected.

Start putting your code in functions. Even for small programs.

from collections import OrderedDict

ODL = [...]

MAXES = {
    ('A', None): 5,
    ('B', None): 5,
    ('C', 'C_1'): 2,
    ('C', 'C_2'): 2,
}

def main():
    for d in select_dicts(ODL):
        print(d)

def select_dicts(ods):
    sort_key = lambda od: (od['date'], od['time'])
    type_key = lambda od: (od['type'], od.get('C_TYPE'))
    available = dict(MAXES)
    selected = []
    for od in sorted(ods, key=sort_key):
        k = type_key(od)
        if available[k] > 0:
            selected.append(od)
            available[k] -= 1
    return selected

if __name__ == '__main__':
    main()