dictionary comprehension with list iteration based on value dimensions

Question

Given

import numpy as np

t={'a':np.random.randint(0,9,[2,3]),'b':np.random.randint(0,9,[2,4])}
l=np.hstack([t.get(k) for k in t.keys()])
l=np.vstack((l, np.random.randint(0,9,[1,7])))

Is there a way to map list l in the above to a dictionary such that the keys map to the keys in dictionary t and values map to the modified list l aligned on columns, same as in t ?

The following for loop works:

t2={}
s=0
e=0
for k in t.keys():
    e=s+t.get(k).shape[1]
    t2[k]=l[:,s:e]
    s=e

but I was wondering if there is a one liner dictionary comprehension equivalent to the above for loop?

Answer 1

Your code depends on the order of the keys. Try to replace for k in t.keys() by for k in ('b', 'a') and you won't have the expected result. Even if the order of elements is guaranted since Python 3.7 (see comment) this might not be a good idea to rely on it.

You could use tuples:

import numpy as np

a = np.random.randint(0,9,[2,3])
b = np.random.randint(0,9,[2,4])

t = [('a', a), ('b', b)]
l = np.hstack([v for _, v in t])

c = np.random.randint(0,9,[1,7])
l = np.vstack((l, c))

To convert your code in a dict comprehension, you can compute the column indices:

import itertools
indices = list(itertools.accumulate((v.shape[1] for _, v in t)))
# [3, 7]

And then produce the arrays:

t2 = {k: l[:,s:s+v.shape[1]] for (k, v), s in zip(t, [0]+indices)}