Select rows of a matrix using binary array

Viewed 471

Supposing that I have a matrix with a format (m x n) and one array with dimension (m) containing booleans. I would like to select and extract just the rows of my matrix in which the corresponding index of the m-dimension array contains a true value. There must be a really simple way to solve this issue which I am not aware of.

A minimal reproducible able example that might help to better explain:

A = np.array([[ 1,  4,  5, 12],
              [-5,  8,  9,  0],
              [-6,  7, 11, 19],
              [13, 15, 16, 19]])

B = np.array([1,0,1,1])

Expected output:

Out[1]: 
array([[ 1,  4,  5, 12],
       [-6,  7, 11, 19],
       [13, 15, 16, 19]])
1 Answers

Cast B to bool, so the indexing is boolean based:

A[B.astype(bool)]

array([[ 1,  4,  5, 12],
       [-6,  7, 11, 19],
       [13, 15, 16, 19]])

Otherwise, being B an array of integers, you'll be performing integer indexing, and will just be indexing on the rows specified by the indices:

A[B]

array([[-5,  8,  9,  0],
       [ 1,  4,  5, 12],
       [-5,  8,  9,  0],
       [-5,  8,  9,  0]])

Find more on boolean indexing here

Related