Question

It’s hard to explain what I’m trying to do with words so here’s an example.

Let’s say we have the following inputs:

In [76]: x
Out[76]: 
0    a
1    a
2    c
3    a
4    b

In [77]: z
Out[77]: ['a', 'b', 'c', 'd', 'e']

I want to get:

In [78]: ii
Out[78]: 
array([[1, 0, 0, 0, 0],
       [1, 0, 0, 0, 0],
       [0, 0, 1, 0, 0],
       [1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0]])

ii is an array of boolean masks which can be applied to z to get back the original x.

My current solution is to write a function which converts z to a list and uses the index method to get the index of the element in z and then generate a row of zeroes except for the index where there is a one. This function gets applied to each row of x to get the desired result.

Answer 1

A first possibility:

>>> choices = np.diag([1]*5)
>>> choices[[z.index(i) for i in x]]

As noted elsewhere, you can change the list comprehension [z.index(i) for i in x] by np.searchsorted(z, x)

>>> choices[np.searchsorted(z, x)]

Note that as suggested in a comment by @seberg, you should use np.eye(len(x)) instead of np.diag([1]*len(x)). The np.eye function directly gives you a 2D array with 1 on the diagonal and 0 elsewhere.