How to create a boolean array where the value is based on an array of indices?

Viewed 519

Let say I have a numpy array A as follows:

A = 
array([[0, 2],
       [1, 2],
       [0, 1]])

I created a boolean array B using np.zeros as follows

B = 
array([[False, False, False],
       [False, False, False],
       [False, False, False]])

Now, I wanted to create an array C, where it gives True value based on the column index of A.

So,

C = 
array([[True, False, True],
       [False, True, True],
       [True, True, False]])
1 Answers

You can do this using some Numpy's relatively advanced indexing techniques:

In [27]: B[np.arange(A.shape[0])[:,None], A] = True                                                                                                                                                  

In [28]: B                                                                                                                                                                                                  
Out[28]: 
array([[ True, False,  True],
       [False,  True,  True],
       [ True,  True, False]])

The np.arange(A.shape[0])[:,None] creates the following array

array([[0],
       [1],
       [2]])

to be used as indices for the first axis of the B array. The [:,None] here is used to transform the one-dimensional range object into a two dimensional array to match with the second axis indices (A array ) which is 2 dimensional.

Related