Select subset of rows of numpy array based on a selection of rows in another array-CodePudding

I have 2 numpy arrays and I wanted to select a subset of rows in one of them based on a conditional subset of the other

arr1 = np.array([[1, 2, 1, 5], [3, 4, 1, 6], [2, 2, 2, 7]])
arr2 = np.array([[2, 2, 1], [2, 3, 0], [2, 1, 1]])

I want to select only those rows in arr1 where the 3rd element in arr2 is 1. In this case, arr2 will look like so: np.array([[2, 2, 1], [2, 1, 1]]) and arr1 will become: np.array([[1, 2, 1, 5], [2, 2, 2, 7]]). Both of them can be assumed to have the same number of rows, but can have different number of columns. How can I achieve this?

CodePudding user response：

In [500]: arr2
Out[500]: 
array([[2, 2, 1],
       [2, 3, 0],
       [2, 1, 1]])

3rd element of each row:

In [502]: arr2[:,2]
Out[502]: array([1, 0, 1])
In [503]: arr2[:,2]==1
Out[503]: array([ True, False,  True])

apply this boolean mask to select rows of arr1:

In [504]: arr1[arr2[:,2]==1]
Out[504]: 
array([[1, 2, 1, 5],
       [2, 2, 2, 7]])