Is there a good comprehension way to access several 2D arrays burried in dicts

user_13816248 0
Viewed 48

So for legacy reasons I have to deal with a data store that has the format

obj.dict1 = {keyA:{key1:(3xN np array),
                   key2:(3xN np array),
                   key3:(3xN np array)}
             keyB:{key1:(3xN np array),
                   key2:(3xN np array),
                   key3:(3xN np array)}
            }

and I am trying to get two different lists quickly:

  1. [ obj.dict1[keyA][key1](0,:), obj.dict1[keyA][key2](0,:), obj.dict1[keyA][key3](0,:) ]

  2. [ obj.dict1[keyA][key1](0,:), obj.dict1[keyB][key1](0,:) ]

So far for both of these I am building a for loop that goes over either the inner or out keys. It works alright but was hoping for something more pythonic with list comprehension

EDIT

The loop for 1. is

keyList = ['key1','key2','key3']
tmpArray = np.zeros((3,N))
for ii in range(3):
    tmpArray[ii,:] = obj.dict1['keyA'][keyList[ii]][0,:]

If it can't be faster or simpler that's fine but this is then inside another for loop that gets repeated quite a few (>50) times.

python
1 Answers

I tried this list comprehension [dict1[key][key2][0,:] for key in dict1 for key2 in dict1[key]] and it ran pretty quickly... in fact, time taken 0.0019943714141845703 seconds for N = 100000

import numpy as np

N = 100000
dict1 ={
    "keyA":{
             "key1": np.random.rand(3,N),
             "key2": np.random.rand(3,N),
             "key3": np.random.rand(3,N)
            },

     "keyB":{
             "key1": np.random.rand(3,N),
             "key2": np.random.rand(3,N),
             "key3": np.random.rand(3,N)
            }
}

temp = []
for key in dict1:
    for key2 in dict1[key]:
        temp.append(dict1[key][key2][0:1])
        
start = time.time()
print([dict1[key][key2][0,:] for key in dict1 for key2 in dict1[key]])
stop = time.time()   
print("time taken",stop-start)
user_115738420
Related