I need to identify which object(image array in this case) a particular index points to from two dictionaries.

dict_count = {1: 400, 2:300, 3:200, 4: 100} 

dict_count says I have 400 of object 1, 300 of object 2 and so on.

dict_object = {1 :[obj1, obj2 ... obj400], 2 :[obj1, obj2 ...obj300], 3 :[obj1, obj2 ...], 4 :[obj1, obj2 ...] 

dict_object contains the objects (max count is mentioned in dict_count)

Now, if I am given an index = 450, I need to get the 50th object under ‘label 2’ from the dict_object.

Existing code (probably not efficient)

Matrix = [[-1 for x in range(sum(dict_count.values()))] for y in range(2)] 
i = 0
index = 450
for k in dict_object.keys():
    for images in dict_object[k]:
        Matrix[0][i] = images
        Matrix[1][i] = k
        i += 1
        
obj = Matrix[index-1][0]
label = MAtrix[index-1][1]

How do I do this efficiently? TIA!

>Solution :

Rather than iterating through indices 1 at a time, iterate through dict_count so you can do it in chunks:

def get_obj(i):
    for k, c in dict_count.items():
        if i < c:
            return dict_object[k][i]
        else:
            i -= c
    raise IndexError(f"index {i} exceeds {sum(dict_count.values())}")