Выбор строк в зависимости от условий для нескольких столбцов

Question

У меня есть массив, и я хочу выбрать строки, где у меня есть некоторые условия для разных столбцов этих строк в Python с помощью NumPy. Для примера рассмотрим этот массив:

test_array = numpy.array(([1,2,3,5],[4,5,6,7],[7,8,9,4])) 

Теперь я хочу, чтобы все строки, в которых столбец 1 является 1 и столбец 4 является 5. Таким образом, выход должен быть [1,2,3,5], только 1 строка в этом случае.

Answer 1

Вы можете использовать следующий подход:

>>> test_array[np.where(np.all(test_array[:,[0,3]]==[1,5],axis=1))] 
array([[1, 2, 3, 5]]) 

Answer 2

[list(x) for x in test_array if x[0]==1 and x[3]==5] 

Это дает желаемый результат:

[[1, 2, 3, 5]] 

Для массива, как этот

test_array=numpy.array(([1,2,3,5],[4,5,6,7],[7,8,9,4],[1,98,76,5])) 

вы бы тогда получить

[[1, 2, 3, 5], [1, 98, 76, 5]] 

Предполагая, что вы называете результат res , вы можете легко получить доступ к результатам. например res[1] будет [1, 98, 76, 5].

Если - по какой-то причине - вы хотели бы получить два числа между 1 и 5 можно использовать

[sl[1:3] for sl in res] 

, который даст вам

[[2, 3], [98, 76]] 

Это то, что вы были находясь в поиске?

Answer 3

Логической индексация/маскировка делает работу хорошо

In [498]: test_array=np.array(([1,2,3,5],[4,5,6,7],[7,8,9,4])) 

In [499]: I = (test_array[:,0]==1) & (test_array[:,3]==5) 

In [500]: I 
Out[500]: array([ True, False, False], dtype=bool) 

In [501]: test_array[I,:] 
Out[501]: array([[1, 2, 3, 5]]) 

При расчете I используйте () свободно поэтому == тестов имеют приоритет над &. Kasra s тест также будет работать: I = np.all(test_array[:,[0,3]]==[1,5],axis=1).

Answer 4

Если вы хотите функцию, чтобы быть критерии для строк, вы можете сделать:

def filter(row): 
    if row[0]==1 and row[3]==5: 
     return True 
    return False 

out_array = test_array[numpy.array([filter(row) for row in test_array])] 

numpy.array([filter(row) for row in test_array]) дает булево массив, соответствующий, какие строки вы хотите, и которые вы не знаете.