Styjun

I have a dictionary in python which looks like this:

mydict=
 "key1": np.array([1, 2, 3]),
 "key2": np.array([4, 5, 6]),
 "key3": np.array([7, 8, 9])
 

Now, I'd like to obtain the dictionary but with just the first the nth entry in each of the value arrays. Something like mydict2=mydict[:][1]. The expected output would be:

mydict2=
 "key1": 2
 "key2": 5
 "key3": 8
 

And I would like to get parts of all arrays simultaneously. For example mydict3=mydict[:][:2]. Here, I expect:

mydict3=
 "key1": np.array([1, 2]),
 "key2": np.array([4, 5]),
 "key3": np.array([7, 8])
 

Obviously, indexing via [:] doesn't work.
How do I achieve this?

What you are describing is a feature mostly offered by numpy, where slicing gives you views. Dictionaries are different beasts, partly because dictionaries are unordered.

You'll have to iterate in a dict comprehension:

mydict2 = k: v[1] for k, v in mydict.items()

and

mydict3 = k: v[:2] for k, v in mydict.items()

I'd simply not store numpy arrays in a dictionary, and use a larger array instead:

fullarray = np.arange(1, 10).reshape(3, 3)

and, at most, use a dictionary to map keys to slices on that array. These slices give you views so altering the full array is reflected in the references in the dictionary:

>>> import numpy as np
>>> fullarray = np.arange(1, 10).reshape(3, 3)
>>> mydict1 = 'key1': fullarray[0, :], 'key2': fullarray[1, :], 'key3': fullarray[2, :]
>>> mydict1
'key1': array([1, 2, 3]), 'key2': array([4, 5, 6]), 'key3': array([7, 8, 9])
>>> fullarray[:, 1]
array([2, 5, 8])
>>> fullarray[:, 1] *= 2
>>> fullarray[:, 1]
array([ 4, 10, 16])
>>> mydict1
'key1': array([1, 4, 3]), 'key2': array([ 4, 10, 6]), 'key3': array([ 7, 16, 9])

but accessing a column is just simpler via the full numpy array, so fullarray[:, 0] and fullarray[:, :2].

Another option is to use structured arrays to produce rows with names:

>>> import numpy.lib.recfunctions as rfn
>>> structured = rfn.unstructured_to_structured(np.arange(1, 10).reshape((3, 3)).T, names=('key1', 'key2', 'key3'))

at which point indexing with a key name gives you arrays:

>>> structured['key1']
array([1, 2, 3]

but you can also index by 'column', including slicing:

>>> structured[0]
(1, 4, 7)
>>> structured[:2]
array([(1, 4, 7), (2, 5, 8)],
 dtype=[('key1', '<i8'), ('key2', '<i8'), ('key3', '<i8')])
>>> structured[:2]['key1']
array([1, 2])
>>> structured[:2]['key2']
array([4, 5])
>>> structured[:2]['key3']
array([7, 8])

Converting an existing dictionary would require stacking the values:

import numpy as np
import numpy.lib.recfunctions as rfn

def dict_to_structured(d):
 return rfn.unstructured_to_structured(
 np.stack(list(mydict.values()), axis=1),
 names=list(mydict)
 )