How to create a list from a dictionary that contains numpy array as value

Question

I have a dictionary like this

{0: array([-6139.66579119, -8102.82498701, -8424.43378713, -8699.96492463,
        -9411.35741859]),

 1: array([ -7679.11144698, -16699.49166421,  -3057.05148494, -10657.0539235 ,
         -3091.04936367]),

 2: array([ -7316.47405724, -15367.98445067,  -6660.88963907,  -9634.54357714,
         -6667.05832509]),

 3: array([-7609.14675848, -9894.14708559, -4040.51364199, -8661.16152946,
        -4363.71589143]),

 4: array([-5068.85919923, -6691.36104136, -6659.66791024, -6666.66570889,
        -5365.35153533]),

 5: array([ -8341.96211464, -13495.42783124,  -4782.52084352, -10355.98002   ,
         -5424.48813488]),

 6: array([ -7740.36341878, -16165.48430318,  -5169.42471878, -12369.79859385,
         -5807.66380805]),

 7: array([-10645.12432969,  -5465.30533986,  -6756.65159092,  -4146.34937333,
         -6765.69595854]),

 8: array([ -7765.04423986, -11679.3889257 ,  -4218.9629257 ,  -6565.64225892,
         -4538.09199979]),

 9: array([-5869.18259848, -7809.21110907, -3272.33611955, -3881.64743889,
        -3275.54657818])}

What I want to do is:

1. compare the first value in each array, in this case, -6139, -7649......and find the max value (-5068), then return the key 4 in a list.

2. compare the second value in each array, -8102, -16699......find the max and return the key , append to the list.

How can I do that?

My code is like this:

def predict(trainingData, testData):

pred = {}

maxLabel = None

prediction=[]

maxValue = -9999999999

pred = postProb(trainingData, testData)


for key, value in pred.items(): 

    for i in range(value.shape[0]):

        for j in range(10):

            if pred[j][i] > maxValue:

                maxValue = pred[key][i]

                maxLabel = key

        prediction.append(maxLabel)

return prediction

pred is the dictionary. It seems that the first loop is not necessary but I need it to get through the elements in the dictionary

Answer 1

You can use numpy array's argmax method to get what you want.

np.array(list(abc.values())).argmax(axis=0)

Out: array([4, 7, 1, 9, 1])

This works only if your keys are consecutive integers like in your example. IF you want a more fool proof method, You could use pandas .

import pandas as pd
df = pd.DataFrame(my_dict)
my_list = list(df.idxmax(axis=1))
print(my_list)

Out: [4, 7, 1, 9, 1]