Получить подмножество строк в матрице numpy на основе значений из столбца другой матрицы

Question

Название выглядит сложным, но проблема не такая уж сложная. У меня есть 2 матрицы: data_X и data_Y. Я должен построить новую матрицу на основе data_X, которая будет состоять из всех строк data_X, где соответствующее значение в столбце column в data_Y не равно someNumber. То же самое для data_Y. Например, здесь находится 5 на 2 data_X матрица и 5 на 1 data_Y матрица, column - 0 и someNumber = -1.

[[ 0.09580361 0.11221975] 
[ 0.71409124 0.24583188] 
[ 0.67346718 0.72550385] 
[ 0.40641294 0.01172211] 
[ 0.89974846 0.70378831]] # data_X 

и data_Y = np.array([[5], [-1], [4], [2], [-1]]).

Результат будет выглядеть так:

[[ 0.09580361 0.11221975] 
[ 0.67346718 0.72550385] 
[ 0.40641294 0.01172211]] 
[5 4 2] 

Это не трудно понять, что это может быть достигнуто за счет следующих действий:

data_x, data_y = [], [] 
for i in xrange(len(data_Y)): 
    if data_Y[i][column] != someNumber: 
     data_y.append(data_Y[i][column]) 
     data_x.append(data_X[i]) 

Но я верю, что есть способ более простой способ (например, 2 или 3 операции numpy), чтобы получить нужные результаты.

Answer 1

Использование boolean indexing -

In [228]: X 
Out[228]: 
array([[ 0.09580361, 0.11221975], 
     [ 0.71409124, 0.24583188], 
     [ 0.67346718, 0.72550385], 
     [ 0.40641294, 0.01172211], 
     [ 0.89974846, 0.70378831]]) 

In [229]: Y 
Out[229]: 
array([[ 5], 
     [-1], 
     [ 4], 
     [ 2], 
     [-1]]) 

In [230]: mask = Y!=-1 # Create mask for boolean indexing 

In [231]: X[mask.ravel()] 
Out[231]: 
array([[ 0.09580361, 0.11221975], 
     [ 0.67346718, 0.72550385], 
     [ 0.40641294, 0.01172211]]) 

In [232]: Y[mask] 
Out[232]: array([5, 4, 2])