Extract values from a numpy array based on another array of 0/1 indices

Question

Given an index array idx that only contains 0 and 1 elements, and 1s represent the sample indices of interest, and a sample array A (A.shape[0] = idx.shape[0]). The objective here is to extract a subset of samples based on the index vector.

In matlab, it is trivial to do:

B = A(idx,:) %assuming A is 2D matrix and idx is a logical vector

How to achieve this in Python in a simple manner?

Answer 1

If your mask array idx has the same shape as your array A, then you should be able to extract elements specified by the mask if you convert idx to a boolean array, using astype.

Demo -

>>> A
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
>>> idx
array([[1, 0, 0, 1, 1],
       [0, 0, 0, 1, 0],
       [1, 0, 0, 1, 1],
       [1, 0, 0, 1, 1],
       [0, 1, 1, 1, 1]])

>>> A[idx.astype(bool)]
array([ 0,  3,  4,  8, 10, 13, 14, 15, 18, 19, 21, 22, 23, 24])

Answer 2

Using the bool operation is equivalent to that logical one in Matlab:

B = A[idx.astype(bool)]