1

Say I have a dictionary:

d = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}

Given a numpy array** of keys, e.g. ['a' 'c' 'e'], is there an easy way to create a numpy array of the corresponding values, e.g. [1 3 5]?

** for reasons of the context I'm using this in, I'd prefer to input/output numpy arrays. Something similar can obviously be done with lists:

keys = ['a', 'c', 'e']
values = []

for k in keys:
    values.append(d[k])

print(values)
[1, 3, 5]

But I am wondering if there is a built-in way to do something like this in numpy.

Subhrajyoti Das
  • 2,685
  • 3
  • 21
  • 36
TY Lim
  • 509
  • 1
  • 3
  • 11

2 Answers2

1

You could do this with a structured array and np.in1d.

import numpy as np

d = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}

# Create structured array (modify formats, names as necessary):
names = ['key', 'data']  
formats = ['U10', 'i4']                                                                        
dtype = dict(names = names, formats = formats)

arr = np.array(list(d.items()), dtype=dtype)

# >>> arr
# array([('a', 1), ('b', 2), ('c', 3), ('d', 4), ('e', 5)],
#       dtype=[('key', '<U10'), ('data', '<i4')])

# use np.in1d to get your keys:

keys = ['a', 'c', 'e']

values = arr[np.in1d(arr['key'], keys)]['data']                                                

# >>> values
# array([1, 3, 5], dtype=int32)
sacuL
  • 49,704
  • 8
  • 81
  • 106
1

If you just require to have numpy arrays as input and output but don't insist on using a numpy function to handle them, you can use numpy.array() on the result of a list comprehension:

import numpy as np

d = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}
keys = np.array(['a', 'c', 'e'])
values = np.array([d[k] for k in keys])

print(values)
LWChris
  • 3,320
  • 1
  • 22
  • 39
  • A list comprehension is a streamlined for loop, but otherwise not much a change from the OP's approach. – hpaulj Nov 23 '19 at 17:19
  • I would still recommend to use list comprehension over append, because it is [much faster](https://stackoverflow.com/q/30245397/1843468). – LWChris Nov 23 '19 at 17:23