SCB
SCB

Reputation: 6149

Numpy broadcast through dictionary

I've got a 2 dimensional numpy array and a dictionary which maps values found in the first column of the array to other values. For example:

>>> x = np.array([[14, 4], [18, 2], [15, 7]])
>>> d = {5: 0, 7: 2, 14: 3, 15: 12, 16: 10, 18: 30}

All keys in d are not guaranteed to be in x though all values in the first column of x will be in d. What I want to do is replace the values in the first column of x with the relevant values in d. Something like:

>>> x[:, 0] = d[x[:, 0]]

So the new array would be:

>>> x
array([[3, 4], [30, 2], [12, 7]])

Of course, this doesn't work because I'm essentially just passing the entire array into the dictionary which wants a key. The best I've come up with is using a for loop:

>>> for i in range(x.shape[0]):
...     x[i, 1] = d[x[i, 1]]

Which of course is very un-numpy and probably not that efficient. My question, is there a "numpy way" of doing something like this?

Upvotes: 2

Views: 1300

Answers (2)

Divakar
Divakar

Reputation: 221564

Here's a Numpythonic solution -

# Extract values and keys
dv = np.array(list(d.values()))
dk = np.array(list(d.keys()))

# Get positions of keys in first column of x and thus change the first column
_,C = np.where(x[:,0][:,None] == dk)
x[:,0] = dv[C]

Sample run -

In [107]: x
Out[107]: 
array([[15,  4],
       [18,  2],
       [14,  7]])

In [108]: d
Out[108]: {16: 10, 18: 3, 5: 0, 7: 2, 14: 12, 15: 30}

In [109]: # Extract values and keys
     ...: dv = np.array(list(d.values()))
     ...: dk = np.array(list(d.keys()))
     ...: 
     ...: # Get positions of keys in first column of x and thus change the first column
     ...: _,C = np.where(x[:,0][:,None] == dk)
     ...: x[:,0] = dv[C]
     ...: 

In [110]: x
Out[110]: 
array([[30,  4],
       [ 3,  2],
       [12,  7]])

Upvotes: 5

Eric
Eric

Reputation: 97571

If your dictionary keys are small and positive:

lookup = np.arange(max(d) + 1)  # or max(x[:,0])
lookup[list(d.keys())] = np.array(d.values())

x[:,0] = lookup[x[:,0]]

Upvotes: 0

Related Questions